"""
v0 script written by Myriam Benisty and Stefano Facchini (Feb 2st 2022)
v1 version adapted by Stefano Facchini and Myriam Benisty (Mar 2nd 2022)
v2 version adapted to include ACA data by Jane Huang and compacted code by Gianni Cataldi (May 19th 2022)
v3 (current version) includes EB alignment by Ryan Loomis and Richard Teague (Aug 2022)


This script is written for CASA 6.2.1.7. 
It can easily be adapted to other versions of CASA, e.g. by modifying the fixvis task to phaseshift task.

The numba package is used for the alignment script, but it is not necessary (it helps with the speed-up)

The scripts are written to be compatible mpicasa.
However, after testingthe exoALMA collaboration decided not to use mpicasa in any of the calibration scripts,
since they realized that differences in the results between casa and mpicasa could arise.

Person in charge for this source:
P. Curone

"""

import os
import numpy as np
import shutil
import matplotlib
#increase the number of figures that can be open before issuing a warning
matplotlib.rcParams.update({'figure.max_open_warning':80})

#path to your local copy of the gihub repository
#(make sure to do a 'git pull' to have the most up-to-date version)
github_path = '/users/pcurone/calibration_scripts'
import sys
sys.path.append(github_path)
import alignment
execfile(os.path.join(github_path,'reduction_utils_exoalma.py'))

prefix = 'PDS_66'

data_folderpath = '/lustre/cv/projects/exoALMA/ALMA_PL_calibrated_data/PDS_66'
TM_names = {'LB':'TM1','SB':'TM2'}
PL_calibrated_vis = {key:os.path.join(data_folderpath,f'{TM_name}/calibrated_final.ms')
                     for key,TM_name in TM_names.items()}

for baseline_key,TM_name in TM_names.items():
    listobs(
        vis=PL_calibrated_vis[baseline_key],
        listfile=f'{prefix}_{TM_name}_calibrated_final.ms.txt',
        overwrite=True,
    )

# System properties.
incl =  30 # deg, from Avenhaus et al. 2018
PA = 189   # deg, from Avenhaus et al. 2018
v_sys = 4     # km/s;  from listobs

# Whether to run tclean in parallel or not.
use_parallel = True

spw_sets = {'SB':['25,27,29,31','93,95,97,99'],
            'LB':['25,27,29,31','92,94,96,98','112,114,116,118','132,134,136,138']}
fields = {'SB':'PDS_66','LB':'PDS_66'}

for baseline_key,vis in PL_calibrated_vis.items():
    split_all_obs(msfile=vis,nametemplate=f'{prefix}_{baseline_key}_EB')


number_of_EBs = {baseline_key:len(spw_set) for baseline_key,spw_set in spw_sets.items()}

for baseline_key,n_EB in number_of_EBs.items():
    for i in range(n_EB):
        vis = f'{prefix}_{baseline_key}_EB{i}.ms'
        listobs(
            vis=vis,
            listfile=f'{vis}.txt',
            overwrite=True,
        )

"""
Check that spws and field numbering is what you expect them to be now from the listobs files
In particular, make sure that there are 4 spws and that there is a single field with
field ID = 0
"""

rest_freq_12CO = 345.79598990e9 #J=3-2
rest_freq_13CO = 330.58796530e9 #J=3-2 (ignoring splitting)
rest_freq_CS   = 342.88285030e9 #J=7-6

data_params_LB = {f'LB{i}': {'vis' : f'{prefix}_LB_EB{i}.ms',
                             'name' : f'LB_EB{i}',
                             'field' : fields['LB'],
                             'line_spws': np.array([0,2,3]), # list of spws containing lines
                             'line_freqs': np.array([rest_freq_13CO,rest_freq_CS,rest_freq_12CO]), #frequencies (Hz) corresponding to line_spws
                             'cont_spws': None,
                             'width_array': None,
                             }
                  for i in range(number_of_EBs['LB'])}

data_params_SB = {f'SB{i}': {'vis' : f'{prefix}_SB_EB{i}.ms',
                             'name' : f'SB_EB{i}',
                             'field' : fields['SB'],
                             'line_spws': np.array([0,2,3]), # list of spws containing lines
                             'line_freqs': np.array([rest_freq_13CO,rest_freq_CS,rest_freq_12CO]), #frequencies (Hz) corresponding to line_spws
                             'cont_spws': None,
                             'width_array': None,
                             }
                  for i in range(number_of_EBs['SB'])}

data_params = data_params_LB.copy()
data_params.update(data_params_SB)

figures_foldername = 'figures'
os.mkdir(figures_foldername)

def get_figures_folderpath(foldername):
    return os.path.join(figures_foldername,foldername)

def make_figures_folder(folderpath):
    if os.path.isdir(folderpath):
        print(f'going to deleted folder {folderpath} and its content')
        valid_answer = 'yes'
        answer = input(f'to confirm, type \'{valid_answer}\': ')
        if answer == valid_answer:
            shutil.rmtree(folderpath)
        else:
            print('aborting')
            return
    os.mkdir(folderpath)
    return folderpath

preselfcal_amp_figures_folder = get_figures_folderpath('1_preselfcal_amp_figures')
make_figures_folder(preselfcal_amp_figures_folder)

#adjust these plot ranges according to your data
plotranges = {'SB':[0,500,0,0.4],#xmin,xmax,ymin,ymax
              'LB':[0,2600,0,0.4]}


for params in data_params.values():
    plot_filename = prefix+'_'+params['name']+'_chan-v-amp_preselfcal.png'
    plotms(vis=params['vis'],
            xaxis='channel',
            yaxis='amplitude',
            field=params['field'],
            ydatacolumn='data',
            avgtime='1e8',
            avgscan=True,
            avgbaseline=True,
            iteraxis='spw',
            #There is a bug in plotms that plots the data wrongly if there is unequal flagging
            #between the polarisations
            #the bug occurs if we plot amp and we average over baselines
            #thus we color by polarization such that
            #we easily identify this issue
            coloraxis='corr', 
            showgui = False,
            exprange='all',
            plotfile=os.path.join(preselfcal_amp_figures_folder,plot_filename)
            )
    baseline_key,_ = params['name'].split('_')
    plotms(vis=params['vis'],
           xaxis='UVdist',
           yaxis='amplitude',
           spw='1',
           field=params['field'],
           ydatacolumn='data',
           avgtime='1e8',
           avgscan=True,
           avgchannel='3840',
           showgui = False,
           plotrange=plotranges[baseline_key],
           plotfile=os.path.join(preselfcal_amp_figures_folder,
                                 prefix+'_'+params['name']+'_uvdist-v-amp_cont_spw_preselfcal.png')
           )


# For PDS 66 everything looks fine, but I'll also try the other plots with single polarizations

#for CQ Tau, the plotms polarisation bug seems indeed there in LB EB1 spw3, so let's plot
#polarizations separately to verify
for corr in ('XX','YY'):
    plot_filename = prefix+'_'+data_params['LB1']['name']\
                         +f'_chan-v-amp_preselfcal_pol{corr}_spw3.png'
    plotms(vis=data_params['LB1']['vis'],
            xaxis='channel',
            yaxis='amplitude',
            field=data_params['LB1']['field'],
            ydatacolumn='data',
            avgtime='1e8',
            avgscan=True,
            avgbaseline=True,
            spw='3',
            correlation=corr,
            showgui = False,
            exprange='all',
            plotfile=os.path.join(preselfcal_amp_figures_folder,plot_filename)
            )
#when plotting polarizations separately, everything looks fine also here

for params in data_params.values():
    flagchannels_string = get_flagchannels(ms_dict=params,output_prefix=prefix,
                                           velocity_range=np.array([-15.,15.])+v_sys)
    avg_cont(ms_dict=params,output_prefix=prefix,flagchannels=flagchannels_string,
             contspws=params['cont_spws'],width_array=params['width_array'])

"""
Double-check that the channels idenfied are at the center of the spws,
due to a potential issue with the data_desc_id key in the ms table of some programs
"""
# Flagchannels input string for LB_EB0: '0:836~3004, 2:794~3043, 3:787~3055'
# Flagchannels input string for LB_EB1: '0:837~3005, 2:794~3043, 3:787~3055'
# Flagchannels input string for LB_EB2: '0:837~3005, 2:794~3043, 3:787~3055'
# Flagchannels input string for LB_EB3: '0:835~3003, 2:796~3045, 3:784~3051'
# Flagchannels input string for SB_EB0: '0:837~3005, 2:794~3043, 3:788~3056'
# Flagchannels input string for SB_EB1: '0:837~3005, 2:794~3043, 3:788~3056'


for baseline_key,n_EB in number_of_EBs.items():
    for i in range(n_EB):
        vis = f'{prefix}_{baseline_key}_EB{i}_initcont.ms'
        listobs(
            vis=vis,
            listfile=f'{vis}.txt',
            overwrite=True,
        )


preselfcal_initcont_amp_folder = get_figures_folderpath(
                                  '2_preselfcal_initcont_amp_figures')
make_figures_folder(preselfcal_initcont_amp_folder)

uv_ranges = {'LB':'125~150m','SB':'125~150m'}

for params in data_params.values():
    vis = prefix+'_'+params['name']+'_initcont.ms'
    baseline_key = params['name'].split('_')[0]
    plotms(vis=vis,
            xaxis='UVdist',
            overwrite=True,
            yaxis='amp',
            coloraxis='spw',
            avgtime='1e8',
            avgscan=True,
            showgui=False,
            plotrange=plotranges[baseline_key],
            plotfile=os.path.join(preselfcal_initcont_amp_folder,
                                  prefix+'_'+params['name']+'_uvdist-v-amp_initcont_preselfcal.png'))

    plotms(vis=vis,
            xaxis='time',
            yaxis='amp',
            avgspw=True,
            uvrange=uv_ranges[baseline_key],
            avgchannel='10000',
            avgbaseline=True,
            coloraxis='corr',
            showgui=False,
            overwrite=True,
            plotfile=os.path.join(preselfcal_initcont_amp_folder,
                                  prefix+'_'+params['name']+'_amp_v_time_initcont_preselfcal.png')
            )
#for SB, we see "waterfall features" (times where the amp suddently sharply decreases)
#we will try to fix this with self-cal

""" Define simple masks and clean scales for imaging """
mask_pa = PA #position angle of mask in degrees
mask_semimajor = 1.5 #semimajor axis of mask in arcsec
mask_semiminor = mask_semimajor*np.cos(incl/180.*np.pi) #semiminor axis of mask in arcsec
mask_ra = '13h22m07.382557s' #taken from listobs
mask_dec = '-69.38.12.65398'

mask = f'ellipse[[{mask_ra},{mask_dec}], [{mask_semimajor:.3f}arcsec,'\
            + f' {mask_semiminor:.3f}arcsec], {mask_pa:.1f}deg]'
# Cellsize: ~beam/6-7
cellsize =  {'LB':'0.015arcsec','SB':'0.040arcsec'}
# Image size: ~primary beam 1.22*lam/A = 32'' with A=12m (19 arcsec)
imsize = {'LB':1200,'SB':400}# primary beam
scales = {'LB':[0,8,15,30,80],'SB':[0,8,15,30]}

noise_annulus = f"annulus[[{mask_ra}, {mask_dec}],['4.arcsec', '6.arcsec']]"
thresholds = {'LB':'0.5mJy','SB':'10mJy'} #clean down to 6 sigma for phase cal

image_png_plot_sizes = [3,10] #sizes in arcsec of the zoomed and overview plots of the pngs

preselfcal_images_png_folder = get_figures_folderpath('3_preselfcal_images')
make_figures_folder(preselfcal_images_png_folder)

for baseline_key,params in zip(('LB','SB'),(data_params_LB,data_params_SB)):
    for EB_key,p in params.items():
        imagename = prefix+'_'+p['name']+'_initcont_image'
        tclean_wrapper(
                      vis=prefix+'_'+p['name']+'_initcont.ms',
                      imagename=imagename,
                      deconvolver='multiscale',
                      scales=scales[baseline_key],
                      mask=mask,
                      threshold=thresholds[baseline_key],
                      cellsize=cellsize[baseline_key],
                      imsize=imsize[baseline_key],
                      parallel=use_parallel,
                      savemodel='modelcolumn'
                    )
        estimate_SNR(f'{imagename}.image',disk_mask=mask,noise_mask=noise_annulus)
        rms = imstat(imagename=f'{imagename}.image',region=noise_annulus)['rms'][0]
        params[EB_key]['rms'] = rms
        generate_image_png(f'{imagename}.image',plot_sizes=image_png_plot_sizes,
                            color_scale_limits=[-3*rms,10*rms],
                            save_folder=preselfcal_images_png_folder)

#PDS_66_LB_EB0_initcont_image.image
#Beam 0.110 arcsec x 0.089 arcsec (22.08 deg)
#Flux inside disk mask: 336.68 mJy
#Peak intensity of source: 15.55 mJy/beam
#rms: 1.56e-01 mJy/beam
#Peak SNR: 99.81

#PDS_66_LB_EB1_initcont_image.image
#Beam 0.138 arcsec x 0.099 arcsec (-4.53 deg)
#Flux inside disk mask: 320.60 mJy
#Peak intensity of source: 26.58 mJy/beam
#rms: 1.60e-01 mJy/beam
#Peak SNR: 165.99

#PDS_66_LB_EB2_initcont_image.image
#Beam 0.139 arcsec x 0.099 arcsec (21.46 deg)
#Flux inside disk mask: 321.91 mJy
#Peak intensity of source: 23.56 mJy/beam
#rms: 1.74e-01 mJy/beam
#Peak SNR: 135.64

#PDS_66_LB_EB3_initcont_image.image
#Beam 0.165 arcsec x 0.106 arcsec (-12.38 deg)
#Flux inside disk mask: 270.61 mJy
#Peak intensity of source: 25.78 mJy/beam
#rms: 1.82e-01 mJy/beam
#Peak SNR: 141.70

#PDS_66_SB_EB0_initcont_image.image
#Beam 0.599 arcsec x 0.496 arcsec (-9.34 deg)
#Flux inside disk mask: 278.95 mJy
#Peak intensity of source: 153.14 mJy/beam
#rms: 1.17e+00 mJy/beam
#Peak SNR: 130.82

#PDS_66_SB_EB1_initcont_image.image
#Beam 0.602 arcsec x 0.492 arcsec (11.84 deg)
#Flux inside disk mask: 286.87 mJy
#Peak intensity of source: 165.98 mJy/beam
#rms: 1.21e+00 mJy/beam
#Peak SNR: 136.78


######
# SELF-CAL INDIVIDUAL EBs
######
""" Self-calibration parameters """
single_EB_contspws = '0~3'
single_EB_spw_mapping = [0,0,0,0]

individual_EB_selfcal_shift_folder = get_figures_folderpath(
                                                   '4_individual_EB_selfcal_and_shift_figures')
make_figures_folder(individual_EB_selfcal_shift_folder)

""" One round of phase-only self-cal """
for params in data_params.values():
    vis = prefix+'_'+params['name']+'_initcont.ms'
    single_EB_p1 = prefix+'_'+params['name']+'_initcont.p1'
    os.system(f'rm -rf {single_EB_p1}')
    gaincal(vis=vis,caltable=single_EB_p1,gaintype='T',spw=single_EB_contspws,
            combine='scan,spw',calmode='p',solint='inf',minsnr=4,minblperant=3)
    """ Print calibration png file """
    plotfilename = prefix+'_'+params['name']\
                        +'_initcont_gain_p1_phase_vs_time.png'
    plotms(single_EB_p1,xaxis='time',yaxis='GainPhase',overwrite=True,showgui=False,
           plotfile=os.path.join(individual_EB_selfcal_shift_folder,plotfilename))
    """ Apply the solutions """
    applycal(vis=vis,spw=single_EB_contspws,spwmap=single_EB_spw_mapping,
             gaintable=[single_EB_p1],interp='linearPD',applymode='calonly',calwt=True)
    split(vis=vis,outputvis=prefix+'_'+params['name']+'_initcont_selfcal.ms',
          datacolumn='corrected')


### Image self-cal'd EBs ###

for baseline_key,params in zip(('LB','SB'),(data_params_LB,data_params_SB)):
    for p in params.values():
        imagename = prefix+'_'+p['name']+'_initcont_selfcal_image'
        tclean_wrapper(vis=prefix+'_'+p['name']+'_initcont_selfcal.ms',
                        imagename=imagename,
                        deconvolver='multiscale',
                        scales=scales[baseline_key],
                        mask=mask,
                        threshold=thresholds[baseline_key],
                        cellsize=cellsize[baseline_key],
                        imsize=imsize[baseline_key],
                        parallel=use_parallel,
                        )
        output_image = f'{imagename}.image'
        estimate_SNR(output_image,disk_mask=mask,noise_mask=noise_annulus)
        rms = p['rms']
        generate_image_png(image=output_image,plot_sizes=image_png_plot_sizes,
                           color_scale_limits=[-3*rms,10*rms],
                           save_folder=individual_EB_selfcal_shift_folder)
'''
2 of 47 solutions flagged due to SNR < 4 in spw=0 at 2021/11/28/13:05:04.8
1 of 47 solutions flagged due to SNR < 4 in spw=0 at 2021/12/02/14:02:33.1
3 of 42 solutions flagged due to SNR < 4 in spw=0 at 2022/07/12/21:52:31.1
'''


for baseline_key,params in zip(('LB','SB'),(data_params_LB,data_params_SB)):
    for p in params.values():
        imagename = prefix+'_'+p['name']+'_initcont_selfcal_image'
        #ratio image, to be compared to the ratio image after alignment
        ref_image = f'{prefix}_{baseline_key}_EB0_initcont_selfcal_image.image'
        ref_rms = params[f'{baseline_key}0']['rms']
        ratio_image = imagename+'.ratio'
        os.system(f'rm -rf {ratio_image}')
        immath(imagename=[ref_image,imagename+'.image'],mode='evalexpr',
               outfile=ratio_image,expr=f'iif(IM0 > 3*{ref_rms}, IM1/IM0, 0)')
        generate_image_png(ratio_image,plot_sizes=[2*mask_semimajor,2*mask_semimajor],
                           color_scale_limits=[0.5,1.5],image_units='ratio',
                           save_folder=individual_EB_selfcal_shift_folder)

#PDS_66_LB_EB0_initcont_selfcal_image.image
#Beam 0.110 arcsec x 0.089 arcsec (22.08 deg)
#Flux inside disk mask: 340.99 mJy
#Peak intensity of source: 15.90 mJy/beam
#rms: 1.34e-01 mJy/beam
#Peak SNR: 118.26

#PDS_66_LB_EB1_initcont_selfcal_image.image
#Beam 0.138 arcsec x 0.099 arcsec (-4.53 deg)
#Flux inside disk mask: 327.32 mJy
#Peak intensity of source: 27.11 mJy/beam
#rms: 9.25e-02 mJy/beam
#Peak SNR: 293.18

#PDS_66_LB_EB2_initcont_selfcal_image.image
#Beam 0.139 arcsec x 0.099 arcsec (21.46 deg)
#Flux inside disk mask: 327.08 mJy
#Peak intensity of source: 24.07 mJy/beam
#rms: 1.34e-01 mJy/beam
#Peak SNR: 179.37

#PDS_66_LB_EB3_initcont_selfcal_image.image
#Beam 0.165 arcsec x 0.106 arcsec (-12.38 deg)
#Flux inside disk mask: 276.08 mJy
#Peak intensity of source: 26.25 mJy/beam
#rms: 1.12e-01 mJy/beam
#Peak SNR: 234.41

#PDS_66_SB_EB0_initcont_selfcal_image.image
#Beam 0.599 arcsec x 0.496 arcsec (-9.34 deg)
#Flux inside disk mask: 299.73 mJy
#Peak intensity of source: 162.86 mJy/beam
#rms: 3.85e-01 mJy/beam
#Peak SNR: 422.48

#PDS_66_SB_EB1_initcont_selfcal_image.image
#Beam 0.602 arcsec x 0.492 arcsec (11.84 deg)
#Flux inside disk mask: 310.17 mJy
#Peak intensity of source: 175.96 mJy/beam
#rms: 4.09e-01 mJy/beam
#Peak SNR: 429.74



######
# ALIGN DATA (go from *initcont_selfcal.ms to *initcont_shift.ms)
######

# Select the LB EB to act as the reference (usually the best SNR one).
# I chose LB EB1 as reference.

reference_for_LB_alignment = f'{prefix}_LB_EB1_initcont_selfcal.ms'
assert 'LB' in reference_for_LB_alignment,\
            'you need to choose an LB EB for alignment of LB'

alignment_offsets = {}

# All the other EBs will be aligned to the reference EB
# We also include the reference EB itself to make sure the coordinate changes are
# copied over.

offset_LB_EBs = ['{}_{}_initcont_selfcal.ms'.format(prefix, params['name'])
                 for params in data_params_LB.values()]

#select the continuum spw with the large bandwidth
continuum_spw_id = 1
# Find the relative offsets and update the phase centers for all offset_EBs.
npix = 1024
cell_size = 0.01
alignment.align_measurement_sets(reference_ms=reference_for_LB_alignment,
                                 align_ms=offset_LB_EBs,npix=npix,cell_size=cell_size,
                                 spwid=continuum_spw_id)

# 2022-10-18 12:42:56     WARN    fixplanets::::casa      The three FIELD table direction reference frame entries for field 0 are not identical in the input data: 0, 21, 21. Will try to continue ...

#New coordinates for PDS_66_LB_EB0_initcont_selfcal.ms
#requires a shift of [-0.02641,0.017629]

#New coordinates for PDS_66_LB_EB1_initcont_selfcal.ms
#no shift, reference MS.

#New coordinates for PDS_66_LB_EB2_initcont_selfcal.ms
#requires a shift of [0.0034318,0.023197]

#New coordinates for PDS_66_LB_EB3_initcont_selfcal.ms
#requires a shift of [-0.012964,0.038713]


#insert offsets from the alignment output
alignment_offsets['LB_EB0'] = [-0.02641,0.017629]
alignment_offsets['LB_EB1'] = [0.0, 0.0]
alignment_offsets['LB_EB2'] = [0.0034318,0.023197]
alignment_offsets['LB_EB3'] = [-0.012964,0.038713]

shifted_LB_EBs = [EB.replace('.ms','_shift.ms') for EB in offset_LB_EBs]

#to check if alignment worked, calculate shift again and verify that shifts are small (i.e.
#a fraction of the cell size):

for shifted_ms in shifted_LB_EBs:
    if shifted_ms == reference_for_LB_alignment.replace('.ms','_shift.ms'):
        #for some reason the fitter fails when computing the offset of an EB to itself,
        #so we skip the ref EB
        continue
    offset = alignment.find_offset(reference_ms=reference_for_LB_alignment,
                                   offset_ms=shifted_ms,npix=npix,cell_size=cell_size,
                                   spwid=continuum_spw_id)
    print(f'#offset for {shifted_ms}: ',offset)

#offset for PDS_66_LB_EB0_initcont_selfcal_shift.ms:  [-4.36116590e-04 -7.22528431e-07]
#offset for PDS_66_LB_EB2_initcont_selfcal_shift.ms:  [2.13896077e-05 3.92162110e-04]
#offset for PDS_66_LB_EB3_initcont_selfcal_shift.ms:  [-0.00033195  0.0001922 ]


# Merge shifted LB EBs for aligning SB EBs
LB_concat_shifted = f'{prefix}_LB_concat_shifted.ms'
os.system(f'rm -rf {LB_concat_shifted}')
concat(vis=shifted_LB_EBs,concatvis=LB_concat_shifted,dirtol='0.1arcsec',
       copypointing=False)

#2022-10-18 13:36:09     WARN    MSConcat::concatenate (file casacore/ms/MSOper/MSConcat.cc, line 996)   Zero or negative scan numbers in MS. May lead to duplicate scan numbers in concatenated MS.
#2022-10-18 13:36:36     WARN    MSConcat::concatenate (file casacore/ms/MSOper/MSConcat.cc, line 996)   Zero or negative scan numbers in MS. May lead to duplicate scan numbers in concatenated MS.
#2022-10-18 13:36:36     WARN    MSConcat::concatenate (file casacore/ms/MSOper/MSConcat.cc, line 996)   Zero or negative scan numbers in MS. May lead to duplicate scan numbers in concatenated MS.


listobs(vis=LB_concat_shifted, listfile=f'{LB_concat_shifted}.txt', overwrite=True)

# Align SB EBs to concat shifted LB EBs
reference_for_SB_alignment = LB_concat_shifted

offset_SB_EBs = ['{}_{}_initcont_selfcal.ms'.format(prefix, params['name'])
                 for params in data_params_SB.values()]

alignment.align_measurement_sets(reference_ms=reference_for_SB_alignment,
                                 align_ms=offset_SB_EBs,npix=npix,cell_size=cell_size,
                                 spwid=continuum_spw_id)

#2022-10-18 13:58:11     WARN    fixplanets::::casa      The three FIELD table direction reference frame entries for field 0 are not identical in the input data: 0, 21, 21. Will try to continue ...

#New coordinates for PDS_66_SB_EB0_initcont_selfcal.ms
#requires a shift of [-0.038057,0.10286]

#New coordinates for PDS_66_SB_EB1_initcont_selfcal.ms
#requires a shift of [-0.015812,0.057497]

alignment_offsets['SB_EB0'] = [-0.038057,0.10286]
alignment_offsets['SB_EB1'] = [-0.015812,0.057497]

shifted_SB_EBs = [EB.replace('.ms','_shift.ms') for EB in offset_SB_EBs]

#check by calculating offset again
for shifted_ms in shifted_SB_EBs:
    offset = alignment.find_offset(reference_ms=reference_for_SB_alignment,
                                   offset_ms=shifted_ms,npix=npix,cell_size=cell_size,
                                   spwid=continuum_spw_id)
    print(f'#offset for {shifted_ms}: ',offset)
#offset for PDS_66_SB_EB0_initcont_selfcal_shift.ms:  [-0.00021185  0.00218858]
#offset for PDS_66_SB_EB1_initcont_selfcal_shift.ms:  [0.00014105 0.001297  ]

# Remove the '_selfcal' part of the names to match the naming convention below.
for shifted_EB in shifted_LB_EBs+shifted_SB_EBs:
    os.system('mv {} {}'.format(shifted_EB, shifted_EB.replace('_selfcal', '')))


""" Check that the images are indeed aligned after the shift """

for baseline_key,params in zip(('LB','SB'),(data_params_LB,data_params_SB)):
    for p in params.values():
        imagename = prefix+'_'+p['name']+'_initcont_shift_image'
        tclean_wrapper(vis=prefix+'_'+p['name']+'_initcont_shift.ms',
                        imagename=imagename,
                        deconvolver='multiscale',
                        scales=scales[baseline_key],
                        mask=mask,
                        threshold=thresholds[baseline_key],
                        cellsize=cellsize[baseline_key],
                        imsize=imsize[baseline_key],
                        parallel=use_parallel,
                      )
        estimate_SNR(f'{imagename}.image',disk_mask=mask,
                     noise_mask=noise_annulus)
        rms = p['rms']
        generate_image_png(f'{imagename}.image',plot_sizes=image_png_plot_sizes,
                           color_scale_limits=[-3*rms,10*rms],
                           save_folder=individual_EB_selfcal_shift_folder)

for baseline_key,params in zip(('LB','SB'),(data_params_LB,data_params_SB)):
    for p in params.values():
        imagename = prefix+'_'+p['name']+'_initcont_shift_image'
        ref_image = f'{prefix}_{baseline_key}_EB0_initcont_shift_image.image'
        ref_rms = params[f'{baseline_key}0']['rms']
        ratio_image = imagename+'.ratio'
        os.system(f'rm -rf {ratio_image}')
        immath(imagename=[ref_image,imagename+'.image'],mode='evalexpr',
               outfile=ratio_image,expr=f'iif(IM0 > 3*{ref_rms}, IM1/IM0, 0)')
        generate_image_png(ratio_image,plot_sizes=[2*mask_semimajor,2*mask_semimajor],
                           color_scale_limits=[0.5,1.5],image_units='ratio',
                           save_folder=individual_EB_selfcal_shift_folder)
'''
2022-10-18 14:13:06     WARN    SDAlgorithmMSClean::takeOneStep (file casa-source/code/synthesis/ImagerObjects/SDAlgorithmMSClean.cc, line 185) MSClean minor cycle stopped at large scale negative or diverging

2022-10-18 14:17:38     WARN    ImageExprCalculator::compute+   [major: 0.1376 arcsec, minor: 0.099437 arcsec, pa: -4.5319 deg]
2022-10-18 14:17:38     WARN    ImageExprCalculator::compute+    vs Axis Lengths: [1, 1]  (NB: Matrix in Row/Column order)
2022-10-18 14:17:38     WARN    ImageExprCalculator::compute+   [major: 0.10967 arcsec, minor: 0.088626 arcsec, pa: 22.082 deg]
2022-10-18 14:17:41     WARN    ImageExprCalculator::compute    image beams are not the same: Axis Lengths: [1, 1]  (NB: Matrix in Row/Column order)
'''

#PDS_66_LB_EB0_initcont_shift_image.image
#Beam 0.110 arcsec x 0.089 arcsec (22.08 deg)
#Flux inside disk mask: 341.37 mJy
#Peak intensity of source: 15.87 mJy/beam
#rms: 1.35e-01 mJy/beam
#Peak SNR: 117.95

#PDS_66_LB_EB1_initcont_shift_image.image
#Beam 0.138 arcsec x 0.099 arcsec (-4.53 deg)
#Flux inside disk mask: 327.44 mJy
#Peak intensity of source: 27.05 mJy/beam
#rms: 9.24e-02 mJy/beam
#Peak SNR: 292.82

#PDS_66_LB_EB2_initcont_shift_image.image
#Beam 0.139 arcsec x 0.099 arcsec (21.46 deg)
#Flux inside disk mask: 327.23 mJy
#Peak intensity of source: 24.07 mJy/beam
#rms: 1.34e-01 mJy/beam
#Peak SNR: 179.48

#PDS_66_LB_EB3_initcont_shift_image.image
#Beam 0.165 arcsec x 0.106 arcsec (-12.38 deg)
#Flux inside disk mask: 275.86 mJy
#Peak intensity of source: 26.22 mJy/beam
#rms: 1.12e-01 mJy/beam
#Peak SNR: 233.55

#PDS_66_SB_EB0_initcont_shift_image.image
#Beam 0.599 arcsec x 0.496 arcsec (-9.34 deg)
#Flux inside disk mask: 299.52 mJy
#Peak intensity of source: 163.20 mJy/beam
#rms: 3.78e-01 mJy/beam
#Peak SNR: 432.01

#PDS_66_SB_EB1_initcont_shift_image.image
#Beam 0.602 arcsec x 0.492 arcsec (11.84 deg)
#Flux inside disk mask: 310.20 mJy
#Peak intensity of source: 176.09 mJy/beam
#rms: 4.12e-01 mJy/beam
#Peak SNR: 427.21

# Also by checking the different images in casaviewer, the different EBs are well aligned

""" Now that everything is aligned, we inspect the flux calibration. """

for params in data_params.values():
    msfile = prefix+'_'+params['name']+'_initcont_shift.ms'
    export_MS(msfile) #export MS contents into Numpy save files

""" Plot deprojected visibility profiles for all data together """
list_npz_files = []
for baseline_key,n_EB in number_of_EBs.items():
    list_npz_files += [f'{prefix}_{baseline_key}_EB{i}_initcont_shift.vis.npz'
                       for i in range(n_EB)]

deprojected_vis_profiles_folder = get_figures_folderpath('5_deprojected_vis_profiles')
make_figures_folder(deprojected_vis_profiles_folder)

plot_deprojected(filelist=list_npz_files,
                 fluxscale=[1.]*(number_of_EBs['LB']+number_of_EBs['SB']),PA=PA,
                 incl=incl,show_err=True,
                 plot_label=os.path.join(deprojected_vis_profiles_folder,
                                         f'{prefix}_flux_scale_EB_preselfcal.png'))

flux_comparison_folder = get_figures_folderpath('6_flux_comparisons')
make_figures_folder(flux_comparison_folder)

# We choose an EB to compare the flux scaling; if possible the best SB EB.
# Since PDS 66 shows decoherence in the SB EBs (waterfalls), I chose as reference 
# the LB EB with the highest SNR (as for the alignment), that is, LB EB1.
 
flux_ref_EB = 'LB_EB1'

for params in data_params.values():
    plot_label = os.path.join(flux_comparison_folder,
                              'flux_comparison_'+params['name']+f'_to_{flux_ref_EB}.png')
    estimate_flux_scale(reference=f'{prefix}_{flux_ref_EB}_initcont_shift.vis.npz',
                        comparison=prefix+'_'+params['name']+'_initcont_shift.vis.npz',
                        incl=incl, PA=PA,plot_label=plot_label)

#The ratio of the fluxes of PDS_66_LB_EB0_initcont_shift.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.92056
#The scaling factor for gencal is 0.959 for your comparison measurement
#The error on the weighted mean ratio is 2.998e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_LB_EB1_initcont_shift.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 1.00000
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 3.147e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_LB_EB2_initcont_shift.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.93186
#The scaling factor for gencal is 0.965 for your comparison measurement
#The error on the weighted mean ratio is 3.775e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_LB_EB3_initcont_shift.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.83836
#The scaling factor for gencal is 0.916 for your comparison measurement
#The error on the weighted mean ratio is 2.592e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SB_EB0_initcont_shift.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.86434
#The scaling factor for gencal is 0.930 for your comparison measurement
#The error on the weighted mean ratio is 2.492e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SB_EB1_initcont_shift.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.92098
#The scaling factor for gencal is 0.960 for your comparison measurement
#The error on the weighted mean ratio is 2.523e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor


"""

SB has overlapping uv ranges with both ACA and LB, so
we select an SB EB as the reference for flux scaling.

In case this source does not have ACA data, simply ignore the ACA steps in the list below.
 
If all SB EBs suffer from decoherence and cannot be used, see below

We correct flux differences >4%, if phase noise looks reasonable

If you need to re-scale fluxes, use command:
rescale_flux(vis=prefix+'_LB_EB0_initcont_shift.ms', gencalparameter=[1.044])

The general procedure is as follows:
1) check flux scaling, and if flux differences are >4%, re-scale the fluxes unless
    there is clear de-coherence (e.g. seen as a systematic decrease of scaling with UV distance)
2) do not scale those EBs with phase decoherence. Proceed with self-cal following the steps below 
   to the point that those EBs are self calibrated
3) check flux scaling again
4) if no flux scaling has been applied in 1), and you now still see a flux offset, then
    apply the flux scaling determined in step 3) to the non-self caled EBs
    and repeat the self-cal

If all SB EBs suffer from phase decoherence and cannot be used as the reference
for flux scaling:
1) concat ACA data and self-cal them in phase
2) concat them to SB data and self-cal them in phase
3) check flux offsets
4) if there are flux offsets, go back to 1), but before re-starting the process apply
    the corrections to the non-self caled ACA and SB data, and re-do steps 1) - 3).
    For this you should add additional code (essentially copy/paste the code from the first
    iteration) that repeats steps 1-3, but saves all output with different filenames. Remember
    to use the right filenames when you continue the script after step 6
5) Check flux offsets of LB EBs to a correct SB EB
6) concat LB data to ACA+SB and continue with script

"""


############ BEGIN of SB self-cal iteration 1 ############

#for phase cal, clean down to 6 sigma

SB_selfcal_folder = get_figures_folderpath('7_selfcal_SB_figures')
make_figures_folder(SB_selfcal_folder)

SB_cont_p0 = prefix+'_SB_contp0'
os.system('rm -rf %s.ms*' %SB_cont_p0)

concat(vis=[f'{prefix}_SB_EB{i}_initcont_shift.ms' for i in range(number_of_EBs['SB'])],
       concatvis=SB_cont_p0+'.ms',dirtol='0.1arcsec',copypointing=False)

listobs(vis=SB_cont_p0+'.ms',listfile=SB_cont_p0+'.ms.txt',overwrite=True)



"""Define new SB mask using new center read from listobs"""
mask_ra = '13h22m07.386700s'
mask_dec = '-69.38.12.66040'

SB_mask= f'ellipse[[{mask_ra},{mask_dec}], [{mask_semimajor:.3f}arcsec,'\
         +f'{mask_semiminor:.3f}arcsec], {mask_pa:.1f}deg]'

noise_annulus_SB = f"annulus[[{mask_ra}, {mask_dec}],['4.arcsec', '6.arcsec']]"

SB_tclean_wrapper_kwargs = {'mask':SB_mask,'deconvolver':'multiscale',
                            'scales':scales['SB'],'savemodel':'modelcolumn',
                            'imsize':imsize['SB'],'cellsize':cellsize['SB'],
                            'robust':0.5,'interactive':False,
                            'parallel':use_parallel,'gridder':'standard'}

#go down to ~6 sigma
tclean_wrapper(vis=SB_cont_p0+'.ms',imagename = SB_cont_p0,
               threshold = '2mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_cont_p0+'.image', disk_mask = SB_mask,noise_mask = noise_annulus_SB)
#PDS_66_SB_contp0.image
#Beam 0.583 arcsec x 0.491 arcsec (2.32 deg)
#Flux inside disk mask: 308.40 mJy
#Peak intensity of source: 168.26 mJy/beam
#rms: 2.37e-01 mJy/beam
#Peak SNR: 709.13

rms_SB = imstat(imagename = SB_cont_p0+'.image',region = noise_annulus_SB)['rms'][0]
generate_image_png(SB_cont_p0+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_SB,10*rms_SB],
                   save_folder=SB_selfcal_folder)

""" Look for references antennas from weblog, and pick the first that are listed, overlapping with all EBs """
# both EBs: DA63, DA52, DA57, DV08

""" Get station numbers """
for ref_ant in ('DA63','DA52', 'DA57', 'DV08'):
    get_station_numbers(SB_cont_p0+'.ms',ref_ant)
#Observation ID 0: DA63@A035
#Observation ID 1: DA63@A035
#Observation ID 0: DA52@A002
#Observation ID 1: DA52@A002
#Observation ID 1: DA57@A001
#Observation ID 0: DV08@A036
#Observation ID 1: DV08@A036
# Why not ID 0 for DA57? Anyway, checking in the weblog, DA57@A001 is present in both SB EBs

SB_contspws = '0~7'
SB_refant   = 'DA63@A035, DA52@A002, DA57@A001, DV08@A036'
SB_spw_mapping = [0,0,0,0,4,4,4,4]

#First round
SB_p1 = prefix+'_SB_p1'
os.system('rm -rf '+SB_p1)
gaincal(vis=SB_cont_p0+'.ms',caltable=SB_p1,
        gaintype='G', spw=SB_contspws,refant=SB_refant, combine='scan,spw',
        calmode='p', solint='inf', minsnr=3., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw', showgui=True)
""" Print calibration png file """
plotms(SB_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_folder,
                             f'{prefix}_SB_p1_gain_phase_vs_time.png')
       )

""" Apply the solutions """
applycal(vis=SB_cont_p0+'.ms',spw=SB_contspws,
         spwmap=SB_spw_mapping,gaintable=[SB_p1], interp='linearPD',
         calwt=True, applymode='calonly')

SB_cont_p1 = SB_cont_p0.replace('p0','p1')
os.system('rm -rf %s.ms*' % SB_cont_p1)
split(vis=SB_cont_p0+'.ms',outputvis=SB_cont_p1+'.ms', datacolumn='corrected')

#again threshold should be ~6sigma
tclean_wrapper(vis=SB_cont_p1+'.ms',imagename = SB_cont_p1,
               threshold = '1.5mJy', **SB_tclean_wrapper_kwargs)
'''
2022-10-19 14:13:09     SEVERE  image::open (file binding/tools/image/image_cmpt.cc, line 4125) Exception Reported: Exception: Unable to open image PDS_66_SB_contp1.ms.
2022-10-19 14:13:09     SEVERE  image::open (file binding/tools/image/image_cmpt.cc, line 4125)+        ... thrown by static casa::ITUPLE casa::ImageFactory::fromFile(const casa6core::String&, casa6core::Bool) at File: casa-source/code/imageanalysis/ImageAnalysis/ImageFactory2.cc, line: 289

Apparently this is bug specific to CASA 6.2 that doesn't like it when the ms and image have the same name
'''
estimate_SNR(SB_cont_p1+'.image',disk_mask=SB_mask,noise_mask=noise_annulus_SB)
#PDS_66_SB_contp1.image
#Beam 0.583 arcsec x 0.491 arcsec (2.32 deg)
#Flux inside disk mask: 308.69 mJy
#Peak intensity of source: 168.37 mJy/beam
#rms: 2.30e-01 mJy/beam
#Peak SNR: 731.07
generate_image_png(SB_cont_p1+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_SB,10*rms_SB],
                   save_folder=SB_selfcal_folder)


#Second round
SB_p2 = SB_p1.replace('p1','p2')
os.system('rm -rf '+SB_p2)
gaincal(vis=SB_cont_p1+'.ms',caltable=SB_p2,
        gaintype='T',spw=SB_contspws,refant=SB_refant, combine='scan, spw',
        calmode='p', solint='360s', minsnr=2., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw', showgui=True)
""" Print calibration png file """
plotms(SB_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_folder,f'{prefix}_SB_p2_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=SB_cont_p1+'.ms',spw=SB_contspws,
         spwmap=SB_spw_mapping, gaintable=[SB_p2], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_cont_p2 = SB_cont_p1.replace('p1','p2')
os.system('rm -rf %s.ms*' % SB_cont_p2)
split(vis=SB_cont_p1+'.ms',outputvis=SB_cont_p2+'.ms',datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_cont_p2+'.ms',imagename = SB_cont_p2,
               threshold = '1.2mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_cont_p2+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#PDS_66_SB_contp2.image
#Beam 0.583 arcsec x 0.491 arcsec (2.32 deg)
#Flux inside disk mask: 309.39 mJy
#Peak intensity of source: 171.05 mJy/beam
#rms: 1.58e-01 mJy/beam
#Peak SNR: 1083.09
generate_image_png(SB_cont_p2+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_SB,10*rms_SB],
                   save_folder=SB_selfcal_folder)



#Third round
SB_p3 = SB_p2.replace('p2','p3')
os.system('rm -rf '+SB_p3)
gaincal(vis=SB_cont_p2+'.ms', caltable=SB_p3,
        gaintype='T', spw=SB_contspws,refant=SB_refant, combine='spw',
        calmode='p', solint='120s', minsnr=2., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_folder,f'{prefix}_SB_p3_gain_phase_vs_time.png')
       )

""" Apply the solutions """
applycal(vis=SB_cont_p2+'.ms', spw=SB_contspws,
         spwmap = SB_spw_mapping, gaintable=[SB_p3], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_cont_p3 = SB_cont_p2.replace('p2','p3')
os.system('rm -rf %s.ms*' % SB_cont_p3)
split(vis=SB_cont_p2+'.ms',outputvis=SB_cont_p3+'.ms',datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_cont_p3+'.ms',imagename = SB_cont_p3,
               threshold = '0.9mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_cont_p3+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#Beam 0.583 arcsec x 0.491 arcsec (2.32 deg)
#Flux inside disk mask: 309.75 mJy
#Peak intensity of source: 172.80 mJy/beam
#rms: 1.31e-01 mJy/beam
#Peak SNR: 1318.21
generate_image_png(SB_cont_p3+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_SB,10*rms_SB],
                   save_folder=SB_selfcal_folder)



#Fourth round
SB_p4 = SB_p3.replace('p3','p4')
os.system('rm -rf '+SB_p4)
gaincal(vis=SB_cont_p3+'.ms', caltable=SB_p4,
        gaintype='T', spw=SB_contspws,refant=SB_refant, combine='spw', calmode='p',
        solint='60s', minsnr=2., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw')

""" Print calibration png file """
plotms(SB_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_folder,f'{prefix}_SB_p4_gain_phase_vs_time.png')
       )

applycal(vis=SB_cont_p3+'.ms', spw=SB_contspws,
         spwmap = SB_spw_mapping, gaintable=[SB_p4], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_cont_p4 = SB_cont_p3.replace('p3','p4')
os.system('rm -rf %s.ms*' % SB_cont_p4)
split(vis=SB_cont_p3+'.ms',outputvis=SB_cont_p4+'.ms',datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_cont_p4+'.ms',imagename = SB_cont_p4,
               threshold = '0.8mJy',**SB_tclean_wrapper_kwargs)
estimate_SNR(SB_cont_p4+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#PDS_66_SB_contp4.image
#Beam 0.583 arcsec x 0.491 arcsec (2.32 deg)
#Flux inside disk mask: 309.90 mJy
#Peak intensity of source: 174.55 mJy/beam
#rms: 1.17e-01 mJy/beam
#Peak SNR: 1486.83
generate_image_png(SB_cont_p4+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_SB,10*rms_SB],
                   save_folder=SB_selfcal_folder)



#Fifth round
SB_p5 = SB_p4.replace('p4','p5')
os.system('rm -rf '+SB_p5)
gaincal(vis=SB_cont_p4+'.ms', caltable=SB_p5,
        gaintype='T', spw=SB_contspws,refant=SB_refant, combine='spw', calmode='p',
        solint='20s', minsnr=3., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_folder,f'{prefix}_SB_p5_gain_phase_vs_time.png')
       )

applycal(vis=SB_cont_p4+'.ms', spw=SB_contspws,
         spwmap = SB_spw_mapping, gaintable=[SB_p5], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_cont_p5 = SB_cont_p4.replace('p4','p5')
os.system('rm -rf %s.ms*' % SB_cont_p5)
split(vis=SB_cont_p4+'.ms',outputvis=SB_cont_p5+'.ms',datacolumn='corrected')


""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_cont_p5+'.ms',imagename = SB_cont_p5,
               threshold = '0.8mJy',**SB_tclean_wrapper_kwargs)
estimate_SNR(SB_cont_p5+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#PDS_66_SB_contp5.image
#Beam 0.583 arcsec x 0.491 arcsec (2.32 deg)
#Flux inside disk mask: 310.34 mJy
#Peak intensity of source: 179.62 mJy/beam
#rms: 9.23e-02 mJy/beam
#Peak SNR: 1946.09
generate_image_png(SB_cont_p5+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_SB,10*rms_SB],
                   save_folder=SB_selfcal_folder)


#Check how SB selfcal improved things
#we do the check for each step of the self-cal to see how things improve at each step
non_self_caled_SB_vis = SB_cont_p0
self_caled_SB_visibilities = {'p1':SB_cont_p1,
                              'p2':SB_cont_p2,
                              'p3':SB_cont_p3,
                              'p4':SB_cont_p4,
                              'p5':SB_cont_p5}
    

SB_EBs = ('EB0','EB1')
SB_EB_spws = ('0,1,2,3','4,5,6,7')

for self_cal_step,self_caled_vis in self_caled_SB_visibilities.items():
    for EB_key,spw in zip(SB_EBs,SB_EB_spws):
        nametemplate = f'{prefix}_SB_{EB_key}_{self_cal_step}_compare_amp_vs_time'
        visibilities = [self_caled_vis+'.ms',non_self_caled_SB_vis+'.ms']
        plot_amp_vs_time_comparison(
                nametemplate=nametemplate,visibilities=visibilities,spw=spw,
                uvrange=uv_ranges['SB'],output_folder=SB_selfcal_folder)


all_SB_visibilities = self_caled_SB_visibilities.copy()
all_SB_visibilities['p0'] = SB_cont_p0

for self_cal_step,vis_name in all_SB_visibilities.items():
    #split out SB EBs
    vis_ms = vis_name+'.ms'
    nametemplate = vis_ms.replace('.ms','_EB')
    split_all_obs(msfile=vis_ms,nametemplate=nametemplate)
    exported_ms = []
    for i in range(number_of_EBs['SB']):
        EB_vis = f'{nametemplate}{i}.ms'
        export_MS(EB_vis) #export MS contents into Numpy save files
        exported_ms.append(EB_vis.replace('.ms','.vis.npz'))
    for i,exp_ms in enumerate(exported_ms):
        png_filename = f'flux_comparison_SB_EB{i}_{self_cal_step}_to_{flux_ref_EB}.png'
        plot_label = os.path.join(SB_selfcal_folder,png_filename)
        estimate_flux_scale(reference=f'{prefix}_{flux_ref_EB}_initcont_shift.vis.npz',
                            comparison=exp_ms,incl=incl,PA=PA,plot_label=plot_label)
    fluxscale = [1.,]*number_of_EBs['SB']
    plot_label = os.path.join(SB_selfcal_folder,
                              f'deprojected_vis_profiles_SB_{self_cal_step}.png')
    plot_deprojected(filelist=exported_ms,fluxscale=fluxscale, PA=PA, incl=incl,
                     show_err=True,plot_label=plot_label)

#The ratio of the fluxes of PDS_66_SB_contp5_EB0.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.92331
#The scaling factor for gencal is 0.961 for your comparison measurement
#The error on the weighted mean ratio is 2.586e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SB_contp5_EB1.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 0.95393
#The scaling factor for gencal is 0.977 for your comparison measurement
#The error on the weighted mean ratio is 2.578e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

'''
Commments:

- The selfcal fixed both the waterfalls and the descending trend in the flux ratios (now the flux ratio vs UVdist taking as reference LB EB1 is constant in UVdist)
- By far, the most notable improvements where performed by the fifth round of selfcal, whereas the improvements of all the other 4 steps were smaller.
- A flux offset >4% with the LB EB1 remains. 

Decision after presenting results so far to the exoALMA meeting:
- due to the decoherence even at LB, making it impossible to discern a correct flux offset (due to the 'downward trend' in the flux ratio),
  proceed with the SB+LB phase-only selfcal and only then check the flux ratios. 
  In case they're still >4%, apply the rescalings and redo all the selfcal (SB and SB+LB) again.
'''


############ END of SB self-cal iteration 1 ############



############ BEGIN of COMBINED SB+LB phase-only self-cal iteration 1 ############

#phase cal down to 6 sigma

LB_selfcal_folder = get_figures_folderpath('8_selfcal_SBLB_figures')
make_figures_folder(LB_selfcal_folder)

""" Merge the SB self-cal'ed ms with LB ms"""
LB_cont_p0 = prefix+'_SBLB_contp0'
os.system('rm -rf %s.ms*' % LB_cont_p0)

#make sure you choose the right SB MS here; it depends on whether you did two
#iterations of SB self-cal to deal with decoherence
#in the case of PDS 66, I am proceeding now without scaling
concat(vis=[SB_cont_p5+'.ms']+[f'{prefix}_LB_EB{i}_initcont_shift.ms' for i
                                      in range(number_of_EBs['LB'])],
            concatvis=LB_cont_p0+'.ms',dirtol='0.1arcsec',copypointing=False)
#I got a warning about possibility of duplicate scan numbers, but looking at the
#listobs output, everything seems ok

listobs(vis=LB_cont_p0+'.ms',listfile=LB_cont_p0+'.ms.txt',overwrite=True)

mask_ra = '13h22m07.386700s'
mask_dec = '-69.38.12.66040'

LB_mask = f'ellipse[[{mask_ra},{mask_dec}], [{mask_semimajor:.3f}arcsec,'\
         +f' {mask_semiminor:.3f}arcsec], {mask_pa:.1f}deg]'
noise_annulus_LB = f"annulus[[{mask_ra}, {mask_dec}],['4.arcsec', '6.arcsec']]"

LB_tclean_wrapper_kwargs = {'mask':LB_mask,'deconvolver':'multiscale',
                            'scales':scales['LB'],'savemodel':'modelcolumn',
                            'imsize':imsize['LB'],'cellsize':cellsize['LB'],
                            'robust':0.5,'interactive':False,
                            'parallel':use_parallel,'gridder':'standard'}


""" clean down to ~6 sigma"""
tclean_wrapper(vis=LB_cont_p0+'.ms', imagename = LB_cont_p0,
               threshold = '0.34mJy',**LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p0+'.image', disk_mask=LB_mask, noise_mask=noise_annulus_LB)
rms_LB = imstat(imagename = LB_cont_p0+'.image', region = noise_annulus_LB)['rms'][0]
generate_image_png(LB_cont_p0+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)
#PDS_66_SBLB_contp0.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 309.44 mJy
#Peak intensity of source: 26.58 mJy/beam
#rms: 5.64e-02 mJy/beam
#Peak SNR: 471.70



""" Look for references antennas from weblog, and pick the first that are listed, overlapping with all EBs """
""" Since we are combined SB EBs to LB EBs, ref antennas for both SB and LB are needed """
# For LB:
# EB0: DA43, DA57, DV04, DV25, DA42, DA58
# EB1: DA43, DV06, DA60, DA58, DV04, DV01
# EB2: DA43, DV06, DA58, DV04, DV01, DA57
# EB3: DV06, DV03, DV02, DV04, DV01, DA58

# For SB, we used:'DA63@A035, DA52@A002, DA57@A001, DV08@A036'
for ref_ant in ('DA43','DV06','DA57', 'DV04', 'DA63','DA52'):
    get_station_numbers(LB_cont_p0+'.ms',ref_ant)
# Bear in mind that here ID0=LB0,ID1=LB1, ID2=LB2, ID3,LB3, ID4=SB0, ID5=SB1
#Observation ID 0: DA43@A035
#Observation ID 1: DA43@A035
#Observation ID 2: DA43@A035
#Observation ID 3: DA43@A073
#Observation ID 4: DA43@A073
#Observation ID 5: DA43@A073
#Observation ID 0: DV06@A024
#Observation ID 1: DV06@A024
#Observation ID 2: DV06@A024
#Observation ID 3: DV06@A024
#Observation ID 4: DV06@A024
#Observation ID 5: DV06@A024
#Observation ID 0: DA57@A043
#Observation ID 1: DA57@A043
#Observation ID 2: DA57@A043
#Observation ID 3: DA57@A104
#Observation ID 5: DA57@A001
#Observation ID 0: DV04@A007
#Observation ID 1: DV04@A007
#Observation ID 2: DV04@A007
#Observation ID 3: DV04@A007
#Observation ID 0: DA63@A123
#Observation ID 1: DA63@A123
#Observation ID 2: DA63@A123
#Observation ID 3: DA63@A116
#Observation ID 4: DA63@A035
#Observation ID 5: DA63@A035
#Observation ID 1: DA52@A089
#Observation ID 2: DA52@A089
#Observation ID 4: DA52@A002
#Observation ID 5: DA52@A002


""" Self-calibration parameters """
LB_contspws = '0~23'
LB_refant   = 'DA63@A035, DA52@A002, DA43@A035, DV06@A024, DA57@A043, DV04@A007'
LB_spw_mapping = [0,0,0,0,4,4,4,4,8,8,8,8,12,12,12,12,16,16,16,16,20,20,20,20]


LB_p1 = prefix+'_SBLB.p1'
os.system('rm -rf '+LB_p1)
#if you get many flagged solutions, change gaintype to 'T'
gaincal(vis=LB_cont_p0+'.ms', caltable=LB_p1, gaintype='G', spw=LB_contspws,
        refant=LB_refant, combine='scan,spw', calmode='p', solint='inf',
        minsnr=3., minblperant=4)
'''
4 of 94 solutions flagged due to SNR < 3 in spw=0 at 2021/11/28/13:05:07.8
4 of 94 solutions flagged due to SNR < 3 in spw=8 at 2021/12/02/14:02:31.6
4 of 84 solutions flagged due to SNR < 3 in spw=20 at 2022/07/12/21:52:31.3
'''

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_folder,f'{prefix}_LB_p1_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_cont_p0+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_p1], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_cont_p1 = prefix+'_SBLB_contp1'
os.system('rm -rf %s.ms*' % LB_cont_p1)
split(vis=LB_cont_p0+'.ms', outputvis=LB_cont_p1+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_cont_p1+'.ms',imagename=LB_cont_p1,threshold='0.33mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p1+'.image', disk_mask=LB_mask, noise_mask=noise_annulus_LB)
#PDS_66_SBLB_contp1.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 308.05 mJy
#Peak intensity of source: 26.60 mJy/beam
#rms: 5.53e-02 mJy/beam
#Peak SNR: 481.12
generate_image_png(LB_cont_p1+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)



""" Second round of phase-only self-cal """
LB_p2 = prefix+'_SBLB.p2'
os.system('rm -rf '+LB_p2)
#solint='360s' resulted in "mismatched frequencies" error, so I slighlty changed it
gaincal(vis=LB_cont_p1+'.ms', caltable=LB_p2, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw,scan', calmode='p', solint='360s',
        minsnr=2., minblperant=4)
'''
~4/47 solutions flagged on average
'''

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw')

""" Print calibration png file """
plotms(LB_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_folder,f'{prefix}_LB_p2_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_cont_p1+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_p2], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_cont_p2 = prefix+'_SBLB_contp2'
os.system('rm -rf %s.ms*' % LB_cont_p2)
split(vis=LB_cont_p1+'.ms', outputvis=LB_cont_p2+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_cont_p2+'.ms', imagename = LB_cont_p2, threshold = '0.22mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p2+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_contp2.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 309.78 mJy
#Peak intensity of source: 28.01 mJy/beam
#rms: 3.63e-02 mJy/beam
#Peak SNR: 770.82
generate_image_png(LB_cont_p2+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)



""" Third round of phase-only self-cal """
"""
Check scan length to decide whether to combine scans. If scans are 2 min, do not combine them.
If they are shorter, combine scans here

For PDS 66, scan lenght varies from 2 to 3 minutes for LB, and from 6 to 7 minutes for SB
"""
LB_p3 = prefix+'_SBLB.p3'
os.system('rm -rf '+LB_p3)
gaincal(vis=LB_cont_p2+'.ms', caltable=LB_p3, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='120s',
        minsnr=2., minblperant=4)
#~6/47 flagged solution on average (always under 10)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_folder,f'{prefix}_LB_p3_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_cont_p2+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_p3], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_cont_p3 = prefix+'_SBLB_contp3'
os.system('rm -rf %s.ms*' % LB_cont_p3)
split(vis=LB_cont_p2+'.ms', outputvis=LB_cont_p3+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_cont_p3+'.ms', imagename = LB_cont_p3, threshold = '0.18mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p3+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_contp3.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 310.12 mJy
#Peak intensity of source: 29.24 mJy/beam
#rms: 3.17e-02 mJy/beam
#Peak SNR: 922.24
generate_image_png(LB_cont_p3+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)



""" Fourth round of phase-only self-cal """
LB_p4 = prefix+'_SBLB.p4'
os.system('rm -rf '+LB_p4)
gaincal(vis=LB_cont_p3+'.ms', caltable=LB_p4, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='60s',
        minsnr=2., minblperant=4)
#~6/47 flagged solution on average (always under 10)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_folder,f'{prefix}_LB_p4_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_cont_p3+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_p4], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_cont_p4 = prefix+'_SBLB_contp4'
os.system('rm -rf %s.ms*' % LB_cont_p4)
split(vis=LB_cont_p3+'.ms', outputvis=LB_cont_p4+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_cont_p4+'.ms', imagename = LB_cont_p4, threshold = '0.18mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p4+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_contp4.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 310.52 mJy
#Peak intensity of source: 29.92 mJy/beam
#rms: 3.16e-02 mJy/beam
#Peak SNR: 945.37
generate_image_png(LB_cont_p4+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)


""" Fifth round of phase-only self-cal """
LB_p5 = prefix+'_SBLB.p5'
os.system('rm -rf '+LB_p5)
gaincal(vis=LB_cont_p4+'.ms', caltable=LB_p5, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='30s',
        minsnr=2., minblperant=4)
#up to 9/40 antennas flagged in all LB data

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_folder,f'{prefix}_LB_p5_gain_phase_vs_time.png'))


applycal(vis=LB_cont_p4+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_p5], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_cont_p5 = prefix+'_SBLB_contp5'
os.system('rm -rf %s.ms*' % LB_cont_p5)
split(vis=LB_cont_p4+'.ms', outputvis=LB_cont_p5+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_cont_p5+'.ms', imagename = LB_cont_p5, threshold = '0.18mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p5+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_contp5.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 310.36 mJy
#Peak intensity of source: 30.44 mJy/beam
#rms: 3.25e-02 mJy/beam
#Peak SNR: 936.75

generate_image_png(LB_cont_p5+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)


""" Sixth round of phase-only self-cal """
LB_p6 = prefix+'_SBLB.p6'
os.system('rm -rf '+LB_p6)
gaincal(vis=LB_cont_p5+'.ms', caltable=LB_p6, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='18s',
        minsnr=2., minblperant=4)
# up to 10/40 antennas flagged in all LB data

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_p6,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_p6,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_folder,f'{prefix}_LB_p6_gain_phase_vs_time.png'))


applycal(vis=LB_cont_p5+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_p6], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_cont_p6 = prefix+'_SBLB_contp6'
os.system('rm -rf %s.ms*' % LB_cont_p6)
split(vis=LB_cont_p5+'.ms', outputvis=LB_cont_p6+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_cont_p6+'.ms', imagename = LB_cont_p6, threshold = '0.19mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_cont_p6+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_contp6.image
#Beam 0.143 arcsec x 0.112 arcsec (-0.86 deg)
#Flux inside disk mask: 310.16 mJy
#Peak intensity of source: 30.71 mJy/beam
#rms: 3.30e-02 mJy/beam
#Peak SNR: 931.17
generate_image_png(LB_cont_p6+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],
                   save_folder=LB_selfcal_folder)

#Now check again the flux scaling

self_caled_LB_visibilities = {'p1':LB_cont_p1,
                              'p2':LB_cont_p2,
                              'p3':LB_cont_p3,
                              'p4':LB_cont_p4,
                              'p5':LB_cont_p5,
                              'p6':LB_cont_p6
                              }

for vis in self_caled_LB_visibilities.values(): 
    listobs(vis=vis+'.ms',listfile=vis+'.ms.txt',overwrite=True)

LB_EBs = ('EB0','EB1','EB2','EB3')
LB_EB_spws = ('0,1,2,3', '4,5,6,7', '8,9,10,11', '20,21,22,23') #fill out by referring to listobs output

for self_cal_step,self_caled_vis in self_caled_LB_visibilities.items():
    for EB_key,spw in zip(LB_EBs,LB_EB_spws):
        nametemplate = f'{prefix}_LB_{EB_key}_{self_cal_step}_compare_amp_vs_time'
        visibilities = [self_caled_vis+'.ms',LB_cont_p0+'.ms']
        plot_amp_vs_time_comparison(
                nametemplate=nametemplate,visibilities=visibilities,spw=spw,
                uvrange=uv_ranges['LB'],output_folder=LB_selfcal_folder)

#set to the EB of the combined SBLB data that corresponds to flux_ref_EB
SBLB_flux_ref_EB = 1 #this is LB EB1

total_number_of_EBs = number_of_EBs['SB'] + number_of_EBs['LB']
for self_cal_step,vis in self_caled_LB_visibilities.items():
    nametemplate = f'{prefix}_SBLB_cont{self_cal_step}_EB'
    split_all_obs(msfile=vis+'.ms',nametemplate=nametemplate)
    for i in range(total_number_of_EBs):
        export_MS(f'{nametemplate}{i}.ms')
    for i in range(total_number_of_EBs):
        reference = f'{nametemplate}{SBLB_flux_ref_EB}.vis.npz'
        output = f'flux_comparison_EB{i}_to_EB{SBLB_flux_ref_EB}'\
                       +f'_SBLB_{self_cal_step}.png'
        plot_label = os.path.join(LB_selfcal_folder,output)
        estimate_flux_scale(reference=reference,
                            comparison=f'{nametemplate}{i}.vis.npz',
                            incl=incl, PA=PA,uvbins = np.arange(40.,300.,20.),
                            plot_label=plot_label)
    filelist = [f'{nametemplate}{i}.vis.npz' for i in range(total_number_of_EBs)]
    fluxscale = [1.,]*total_number_of_EBs
    plot_label = os.path.join(LB_selfcal_folder,
                              f'deprojected_vis_profiles_SBLB_{self_cal_step}.png')
    plot_deprojected(filelist=filelist,fluxscale=fluxscale, PA=PA, incl=incl,
                     show_err=True,plot_label=plot_label)

# Redo the flux comparison images without the uvbins parameters, to have clearer plots
for self_cal_step,vis in self_caled_LB_visibilities.items():
    nametemplate = f'{prefix}_SBLB_cont{self_cal_step}_EB'
    for i in range(total_number_of_EBs):
        reference = f'{nametemplate}{SBLB_flux_ref_EB}.vis.npz'
        output = f'flux_comparison_2_EB{i}_to_EB{SBLB_flux_ref_EB}'\
                       +f'_SBLB_{self_cal_step}.png'
        plot_label = os.path.join(LB_selfcal_folder,output)
        estimate_flux_scale(reference=reference,
                            comparison=f'{nametemplate}{i}.vis.npz',
                            incl=incl, PA=PA,
                            plot_label=plot_label)

#The ratio of the fluxes of PDS_66_SBLB_contp6_EB0.vis.npz to
#PDS_66_SBLB_contp6_EB1.vis.npz is 1.04127
#The scaling factor for gencal is 1.020 for your comparison measurement
#The error on the weighted mean ratio is 3.062e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_contp6_EB1.vis.npz to
#PDS_66_SBLB_contp6_EB1.vis.npz is 1.00000
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 3.026e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_contp6_EB2.vis.npz to
#PDS_66_SBLB_contp6_EB1.vis.npz is 1.01705
#The scaling factor for gencal is 1.008 for your comparison measurement
#The error on the weighted mean ratio is 3.733e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_contp6_EB3.vis.npz to
#PDS_66_SBLB_contp6_EB1.vis.npz is 0.83222
#The scaling factor for gencal is 0.912 for your comparison measurement
#The error on the weighted mean ratio is 2.481e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_contp6_EB4.vis.npz to
#PDS_66_SBLB_contp6_EB1.vis.npz is 0.89808
#The scaling factor for gencal is 0.948 for your comparison measurement
#The error on the weighted mean ratio is 2.458e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_contp6_EB5.vis.npz to
#PDS_66_SBLB_contp6_EB1.vis.npz is 0.92549
#The scaling factor for gencal is 0.962 for your comparison measurement
#The error on the weighted mean ratio is 2.443e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

'''
Comments: 
- Phase self-cal fixed the decoherence in the LB EBs, seen as a 'downward trend' in the flux ratio plots.
- The best flux comparison plots are from the sixth round p6 of phase self-cal, even though after p4 the noise in the image increased a tiny bit (from 3.16e-2 mJy to 3.30e-2 mJy). 
  So, I'll consider p6 as final step of the phase self-cal
- Looking at the flux ratios, only LB EB1 (the reference) and LB EB2 don't have to be rescaled.
- So, now I apply rescale_flux to LB EB0, LB EB3, SB EB0, and SB EB1 and then I re-apply all the SB and SB+LB phase calibration.
'''

############ END of COMBINED SB+LB phase-only self-cal iteration 1 ############


# Rescaling flux of the non-self-caled SB and LB EBs

for baseline_key,n_EB in number_of_EBs.items():
    for i in range(n_EB):
        vis = f'{prefix}_{baseline_key}_EB{i}_initcont_shift.ms'
        listobs(
            vis=vis,
            listfile=f'{vis}.txt',
            overwrite=True,
        )

# In the concatenated SBLB file:
# EB4 = SB EB0
# EB5 = SB EB1
# EB0 = LB EB0
# EB1 = LB EB1
# EB2 = LB EB2
# EB3 = LB EB3

rescale_flux(vis=prefix+'_SB_EB0_initcont_shift.ms', gencalparameter=[0.948])
rescale_flux(vis=prefix+'_SB_EB1_initcont_shift.ms', gencalparameter=[0.962])
rescale_flux(vis=prefix+'_LB_EB0_initcont_shift.ms', gencalparameter=[1.020])
rescale_flux(vis=prefix+'_LB_EB3_initcont_shift.ms', gencalparameter=[0.912])
#Splitting out rescaled values into new MS: PDS_66_SB_EB0_initcont_shift_rescaled.ms
#Splitting out rescaled values into new MS: PDS_66_SB_EB1_initcont_shift_rescaled.ms
#Splitting out rescaled values into new MS: PDS_66_LB_EB0_initcont_shift_rescaled.ms
#Splitting out rescaled values into new MS: PDS_66_LB_EB3_initcont_shift_rescaled.ms

listobs(vis=prefix+'_SB_EB0_initcont_shift_rescaled.ms',listfile=prefix+'_SB_EB0_initcont_shift_rescaled.ms'+'.txt',overwrite=True)
listobs(vis=prefix+'_SB_EB1_initcont_shift_rescaled.ms',listfile=prefix+'_SB_EB1_initcont_shift_rescaled.ms'+'.txt',overwrite=True)
listobs(vis=prefix+'_LB_EB0_initcont_shift_rescaled.ms',listfile=prefix+'_LB_EB0_initcont_shift_rescaled.ms'+'.txt',overwrite=True)
listobs(vis=prefix+'_LB_EB3_initcont_shift_rescaled.ms',listfile=prefix+'_LB_EB3_initcont_shift_rescaled.ms'+'.txt',overwrite=True)


############ BEGIN of SB self-cal iteration 2 ############

#reminder: clean down to 6 sigma at each step

SB_selfcal_iteration2_folder = get_figures_folderpath('7.1_selfcal_SB_iteration2_figures')
make_figures_folder(SB_selfcal_iteration2_folder)

SB_iteration2_cont_p0 = prefix+'_SB_iteration2_contp0'
os.system('rm -rf %s.ms*' %SB_iteration2_cont_p0)

concat(vis=[f'{prefix}_SB_EB0_initcont_shift_rescaled.ms',f'{prefix}_SB_EB1_initcont_shift_rescaled.ms',],
       concatvis=SB_iteration2_cont_p0+'.ms',
       dirtol='0.1arcsec',copypointing=False)

listobs(vis=SB_iteration2_cont_p0+'.ms',listfile=SB_iteration2_cont_p0+'.ms.txt',overwrite=True)

mask_ra = '13h22m07.386700s'
mask_dec = '-69.38.12.66040'

SB_mask= f'ellipse[[{mask_ra},{mask_dec}], [{mask_semimajor:.3f}arcsec,'\
         +f'{mask_semiminor:.3f}arcsec], {mask_pa:.1f}deg]'

noise_annulus_SB = f"annulus[[{mask_ra}, {mask_dec}],['4.arcsec', '6.arcsec']]"

SB_tclean_wrapper_kwargs = {'mask':SB_mask,'deconvolver':'multiscale',
                            'scales':scales['SB'],'savemodel':'modelcolumn',
                            'imsize':imsize['SB'],'cellsize':cellsize['SB'],
                            'robust':0.5,'interactive':False,
                            'parallel':use_parallel,'gridder':'standard'}

#clean down to 6 sigma
tclean_wrapper(vis=SB_iteration2_cont_p0+'.ms',imagename = SB_iteration2_cont_p0,
               threshold = '1.5mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_iteration2_cont_p0+'.image',disk_mask=SB_mask,noise_mask=noise_annulus_SB)
#PDS_66_SB_iteration2_contp0.image
#Beam 0.583 arcsec x 0.491 arcsec (2.62 deg)
#Flux inside disk mask: 338.22 mJy
#Peak intensity of source: 184.43 mJy/beam
#rms: 2.47e-01 mJy/beam
#Peak SNR: 745.39


rms_iteration2_SB = imstat(imagename=SB_iteration2_cont_p0+'.image',
                           region=noise_annulus_SB)['rms'][0]
generate_image_png(SB_iteration2_cont_p0+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_SB,10*rms_iteration2_SB],
                   save_folder=SB_selfcal_iteration2_folder)


#First round
SB_iteration2_p1 = prefix+'_SB_iteration2.p1'
os.system('rm -rf '+SB_iteration2_p1)
gaincal(vis=SB_iteration2_cont_p0+'.ms',caltable=SB_iteration2_p1,
        gaintype='G', spw=SB_contspws,refant=SB_refant, combine='scan,spw',
        calmode='p', solint='inf', minsnr=3., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_iteration2_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_iteration2_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_iteration2_folder,
                             f'{prefix}_SB_iteration2_gain_p1_phase_vs_time.png')
       )

""" Apply the solutions """
applycal(vis=SB_iteration2_cont_p0+'.ms',spw=SB_contspws,
         spwmap=SB_spw_mapping,gaintable=[SB_iteration2_p1], interp='linearPD',
         calwt=True, applymode='calonly')

SB_iteration2_cont_p1 = SB_iteration2_cont_p0.replace('p0','p1')
os.system('rm -rf %s.ms*' % SB_iteration2_cont_p1)
split(vis=SB_iteration2_cont_p0+'.ms',outputvis=SB_iteration2_cont_p1+'.ms',
      datacolumn='corrected')

tclean_wrapper(vis=SB_iteration2_cont_p1+'.ms',imagename = SB_iteration2_cont_p1,
               threshold = '1.4mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_iteration2_cont_p1+'.image',disk_mask=SB_mask,noise_mask=noise_annulus_SB)
#PDS_66_SB_iteration2_contp1.image
#Beam 0.583 arcsec x 0.491 arcsec (2.62 deg)
#Flux inside disk mask: 338.30 mJy
#Peak intensity of source: 184.50 mJy/beam
#rms: 2.44e-01 mJy/beam
#Peak SNR: 757.55
generate_image_png(SB_iteration2_cont_p1+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_SB,10*rms_iteration2_SB],
                   save_folder=SB_selfcal_iteration2_folder)



#Second round
SB_iteration2_p2 = SB_iteration2_p1.replace('p1','p2')
os.system('rm -rf '+SB_iteration2_p2)
gaincal(vis=SB_iteration2_cont_p1+'.ms',caltable=SB_iteration2_p2,
        gaintype='T',spw=SB_contspws,refant=SB_refant, combine='scan,spw',
        calmode='p', solint='360s',minsnr=2.,minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_iteration2_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_iteration2_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_iteration2_folder,
                             f'{prefix}_SB_iteration2_gain_p2_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=SB_iteration2_cont_p1+'.ms',spw=SB_contspws,
         spwmap=SB_spw_mapping, gaintable=[SB_iteration2_p2], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_iteration2_cont_p2 = SB_iteration2_cont_p1.replace('p1','p2')
os.system('rm -rf %s.ms*' % SB_iteration2_cont_p2)
split(vis=SB_iteration2_cont_p1+'.ms',outputvis=SB_iteration2_cont_p2+'.ms',
      datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_iteration2_cont_p2+'.ms',imagename = SB_iteration2_cont_p2,
               threshold = '1mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_iteration2_cont_p2+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#PDS_66_SB_iteration2_contp2.image
#Beam 0.583 arcsec x 0.491 arcsec (2.62 deg)
#Flux inside disk mask: 339.04 mJy
#Peak intensity of source: 187.44 mJy/beam
#rms: 1.57e-01 mJy/beam
#Peak SNR: 1196.12
generate_image_png(SB_iteration2_cont_p2+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_SB,10*rms_iteration2_SB],
                   save_folder=SB_selfcal_iteration2_folder)


#Third round
SB_iteration2_p3 = SB_iteration2_p2.replace('p2','p3')
os.system('rm -rf '+SB_iteration2_p3)
gaincal(vis=SB_iteration2_cont_p2+'.ms', caltable=SB_iteration2_p3,
        gaintype='T', spw=SB_contspws,refant=SB_refant, combine='spw',
        calmode='p', solint='120s', minsnr=2., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_iteration2_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_iteration2_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_iteration2_folder,
                             f'{prefix}_SB_iteration2_gain_p3_phase_vs_time.png')
       )

""" Apply the solutions """
applycal(vis=SB_iteration2_cont_p2+'.ms', spw=SB_contspws,
         spwmap = SB_spw_mapping, gaintable=[SB_iteration2_p3], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_iteration2_cont_p3 = SB_iteration2_cont_p2.replace('p2','p3')
os.system('rm -rf %s.ms*' % SB_iteration2_cont_p3)
split(vis=SB_iteration2_cont_p2+'.ms',outputvis=SB_iteration2_cont_p3+'.ms',
      datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_iteration2_cont_p3+'.ms',imagename = SB_iteration2_cont_p3,
               threshold = '0.9mJy', **SB_tclean_wrapper_kwargs)
estimate_SNR(SB_iteration2_cont_p3+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#PDS_66_SB_iteration2_contp3.image
#Beam 0.583 arcsec x 0.491 arcsec (2.62 deg)
#Flux inside disk mask: 339.32 mJy
#Peak intensity of source: 189.32 mJy/beam
#rms: 1.27e-01 mJy/beam
#Peak SNR: 1490.49
generate_image_png(SB_iteration2_cont_p3+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_SB,10*rms_iteration2_SB],
                   save_folder=SB_selfcal_iteration2_folder)



#Fourth round
SB_iteration2_p4 = SB_iteration2_p3.replace('p3','p4')
os.system('rm -rf '+SB_iteration2_p4)
gaincal(vis=SB_iteration2_cont_p3+'.ms', caltable=SB_iteration2_p4,
        gaintype='T', spw=SB_contspws,refant=SB_refant, combine='spw', calmode='p',
        solint='60s', minsnr=2., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_iteration2_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_iteration2_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_iteration2_folder,
                             f'{prefix}_SB_iteration2_gain_p4_phase_vs_time.png')
       )

applycal(vis=SB_iteration2_cont_p3+'.ms', spw=SB_contspws,
         spwmap = SB_spw_mapping, gaintable=[SB_iteration2_p4], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_iteration2_cont_p4 = SB_iteration2_cont_p3.replace('p3','p4')
os.system('rm -rf %s.ms*' % SB_iteration2_cont_p4)
split(vis=SB_iteration2_cont_p3+'.ms',outputvis=SB_iteration2_cont_p4+'.ms',
      datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_iteration2_cont_p4+'.ms',imagename = SB_iteration2_cont_p4,
               threshold = '0.6mJy',**SB_tclean_wrapper_kwargs)
estimate_SNR(SB_iteration2_cont_p4+'.image', disk_mask=SB_mask,
             noise_mask=noise_annulus_SB)
#PDS_66_SB_iteration2_contp4.image
#Beam 0.583 arcsec x 0.491 arcsec (2.62 deg)
#Flux inside disk mask: 339.51 mJy
#Peak intensity of source: 191.20 mJy/beam
#rms: 1.10e-01 mJy/beam
#Peak SNR: 1737.56
generate_image_png(SB_iteration2_cont_p4+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_SB,10*rms_iteration2_SB],
                   save_folder=SB_selfcal_iteration2_folder)



#Fifth round
SB_iteration2_p5 = SB_iteration2_p4.replace('p4','p5')
os.system('rm -rf '+SB_iteration2_p5)
gaincal(vis=SB_iteration2_cont_p4+'.ms', caltable=SB_iteration2_p5,
        gaintype='T', spw=SB_contspws,refant=SB_refant, combine='spw', calmode='p',
        solint='20s', minsnr=3., minblperant=4)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(SB_iteration2_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(SB_iteration2_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(SB_selfcal_iteration2_folder,
                             f'{prefix}_SB_iteration2_gain_p5_phase_vs_time.png')
       )

applycal(vis=SB_iteration2_cont_p4+'.ms', spw=SB_contspws,
         spwmap = SB_spw_mapping, gaintable=[SB_iteration2_p5], interp='linearPD',
         calwt=True, applymode='calonly')

""" Split off a corrected MS """
SB_iteration2_cont_p5 = SB_iteration2_cont_p4.replace('p4','p5')
os.system('rm -rf %s.ms*' % SB_iteration2_cont_p5)
split(vis=SB_iteration2_cont_p4+'.ms',outputvis=SB_iteration2_cont_p5+'.ms',
      datacolumn='corrected')


""" Image the results; check the resulting map """
tclean_wrapper(vis=SB_iteration2_cont_p5+'.ms',imagename = SB_iteration2_cont_p5,
               threshold = '0.6mJy',**SB_tclean_wrapper_kwargs)
estimate_SNR(SB_iteration2_cont_p5+'.image', disk_mask = SB_mask,
             noise_mask = noise_annulus_SB)
#PDS_66_SB_iteration2_contp5.image
#Beam 0.583 arcsec x 0.491 arcsec (2.62 deg)
#Flux inside disk mask: 339.93 mJy
#Peak intensity of source: 196.72 mJy/beam
#rms: 8.84e-02 mJy/beam
#Peak SNR: 2225.93
generate_image_png(SB_iteration2_cont_p5+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_SB,10*rms_iteration2_SB],
                   save_folder=SB_selfcal_iteration2_folder)



#Check again how SB selfcal improved things
#we do the check for each step of the self-cal to see how things improve at each step

non_self_caled_iteration2_SB_vis = SB_iteration2_cont_p0
self_caled_SB_iteration2_visibilities = {'p1':SB_iteration2_cont_p1,
                                         'p2':SB_iteration2_cont_p2,
                                         'p3':SB_iteration2_cont_p3,
                                         'p4':SB_iteration2_cont_p4,
                                         'p5':SB_iteration2_cont_p5}

SB_EBs = ('EB0','EB1')
SB_EB_spws = ('0,1,2,3','4,5,6,7')

for self_cal_step,self_caled_vis in self_caled_SB_iteration2_visibilities.items():
    for EB_key,spw in zip(SB_EBs,SB_EB_spws):
        nametemplate = f'{prefix}_SB_iteration2_{EB_key}_{self_cal_step}_compare'\
                        +'_amp_vs_time'
        visibilities=[self_caled_vis+'.ms',non_self_caled_iteration2_SB_vis+'.ms']
        plot_amp_vs_time_comparison(
                nametemplate=nametemplate,visibilities=visibilities,spw=spw,
                uvrange=uv_ranges['SB'],output_folder=SB_selfcal_iteration2_folder)


all_SB_iteration2_visibilities = self_caled_SB_iteration2_visibilities.copy()
all_SB_iteration2_visibilities['p0'] = SB_iteration2_cont_p0

for self_cal_step,vis_name in all_SB_iteration2_visibilities.items():
    #split out SB EBs
    vis_ms = vis_name+'.ms'
    nametemplate = vis_ms.replace('.ms','_EB')
    split_all_obs(msfile=vis_ms,nametemplate=nametemplate)
    exported_ms = []
    for i in range(number_of_EBs['SB']):
        EB_vis = f'{nametemplate}{i}.ms'
        export_MS(EB_vis) #export MS contents into Numpy save files
        exported_ms.append(EB_vis.replace('.ms','.vis.npz'))
    for i,exp_ms in enumerate(exported_ms):
        png_filename = f'iteration2_flux_comparison_SB_EB{i}_{self_cal_step}'\
                       +f'_to_{flux_ref_EB}.png'
        plot_label = os.path.join(SB_selfcal_iteration2_folder,png_filename)
        estimate_flux_scale(reference=f'{prefix}_{flux_ref_EB}_initcont_shift.vis.npz',
                            comparison=exp_ms,incl=incl,PA=PA,plot_label=plot_label)
    fluxscale = [1.,]*number_of_EBs['SB']
    plot_label = os.path.join(SB_selfcal_iteration2_folder,
                              f'deprojected_vis_profiles_SB_iteration2_{self_cal_step}.png')
    plot_deprojected(filelist=exported_ms,fluxscale=fluxscale, PA=PA, incl=incl,
                     show_err=True,plot_label=plot_label)
#The ratio of the fluxes of PDS_66_SB_iteration2_contp5_EB0.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 1.02739
#The scaling factor for gencal is 1.014 for your comparison measurement
#The error on the weighted mean ratio is 2.877e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SB_iteration2_contp5_EB1.vis.npz to
#PDS_66_LB_EB1_initcont_shift.vis.npz is 1.03079
#The scaling factor for gencal is 1.015 for your comparison measurement
#The error on the weighted mean ratio is 2.786e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor



############ END of SB self-cal iteration 2 ############



############ BEGIN of COMBINED SB+LB phase-only self-cal iteration 2 ############


#phase cal down to 6 sigma, amp cal down to 1 sigma

LB_selfcal_iteration2_folder = get_figures_folderpath('8.1_selfcal_SBLB_iteration2_figures')
make_figures_folder(LB_selfcal_iteration2_folder)

""" Merge the SB self-cal'ed ms with LB ms"""
LB_iteration2_cont_p0 = prefix+'_SBLB_iteration2_contp0'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p0)

# concat the right rescaled files
concat(vis=[SB_iteration2_cont_p5+'.ms', 
            f'{prefix}_LB_EB0_initcont_shift_rescaled.ms',
            f'{prefix}_LB_EB1_initcont_shift.ms',
            f'{prefix}_LB_EB2_initcont_shift.ms',
            f'{prefix}_LB_EB3_initcont_shift_rescaled.ms'],
            concatvis=LB_iteration2_cont_p0+'.ms',
            dirtol='0.1arcsec',copypointing=False)

listobs(vis=LB_iteration2_cont_p0+'.ms',listfile=LB_iteration2_cont_p0+'.ms.txt',overwrite=True)


mask_ra = '13h22m07.386700s'
mask_dec = '-69.38.12.66040'

LB_mask = f'ellipse[[{mask_ra},{mask_dec}], [{mask_semimajor:.3f}arcsec,'\
         +f' {mask_semiminor:.3f}arcsec], {mask_pa:.1f}deg]'
noise_annulus_LB = f"annulus[[{mask_ra}, {mask_dec}],['4.arcsec', '6.arcsec']]"

LB_tclean_wrapper_kwargs = {'mask':LB_mask,'deconvolver':'multiscale',
                            'scales':scales['LB'],'savemodel':'modelcolumn',
                            'imsize':imsize['LB'],'cellsize':cellsize['LB'],
                            'robust':0.5,'interactive':False,
                            'parallel':use_parallel,'gridder':'standard'}


""" clean down to ~6 sigma"""
tclean_wrapper(vis=LB_iteration2_cont_p0+'.ms', imagename = LB_iteration2_cont_p0,
               threshold = '0.34mJy',**LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p0+'.image', disk_mask=LB_mask, noise_mask=noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp0.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 335.78 mJy
#Peak intensity of source: 26.90 mJy/beam
#rms: 6.04e-02 mJy/beam
#Peak SNR: 445.46
rms_iteration2_LB = imstat(imagename = LB_iteration2_cont_p0+'.image', region = noise_annulus_LB)['rms'][0]
generate_image_png(LB_iteration2_cont_p0+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)


""" First round of phase-only self-cal """
LB_iteration2_p1 = prefix+'_SBLB_iteration2.p1'
os.system('rm -rf '+LB_iteration2_p1)
#if you get many flagged solutions, change gaintype to 'T'
gaincal(vis=LB_iteration2_cont_p0+'.ms', caltable=LB_iteration2_p1, gaintype='G', spw=LB_contspws,
        refant=LB_refant, combine='scan,spw', calmode='p', solint='inf',
        minsnr=3., minblperant=4)
'''
2 of 94 solutions flagged due to SNR < 3 in spw=0 at 2021/11/28/13:05:07.5
4 of 94 solutions flagged due to SNR < 3 in spw=8 at 2021/12/02/14:02:31.7
3 of 84 solutions flagged due to SNR < 3 in spw=20 at 2022/07/12/21:52:31.0
'''

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_iteration2_p1,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_p1_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_iteration2_cont_p0+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_p1], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_p1 = prefix+'_SBLB_iteration2_contp1'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p1)
split(vis=LB_iteration2_cont_p0+'.ms', outputvis=LB_iteration2_cont_p1+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_p1+'.ms',imagename=LB_iteration2_cont_p1,threshold='0.36mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p1+'.image', disk_mask=LB_mask, noise_mask=noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp1.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 334.89 mJy
#Peak intensity of source: 26.93 mJy/beam
#rms: 5.94e-02 mJy/beam
#Peak SNR: 453.67
generate_image_png(LB_iteration2_cont_p1+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)


""" Second round of phase-only self-cal """
LB_iteration2_p2 = prefix+'_SBLB_iteration2.p2'
os.system('rm -rf '+LB_iteration2_p2)
gaincal(vis=LB_iteration2_cont_p1+'.ms', caltable=LB_iteration2_p2, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw,scan', calmode='p', solint='360s',
        minsnr=2., minblperant=4)
'''
~5/47 solutions flagged on average (up to 9)
'''

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw')

""" Print calibration png file """
plotms(LB_iteration2_p2,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_p2_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_iteration2_cont_p1+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_p2], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_p2 = prefix+'_SBLB_iteration2_contp2'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p2)
split(vis=LB_iteration2_cont_p1+'.ms', outputvis=LB_iteration2_cont_p2+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_p2+'.ms', imagename = LB_iteration2_cont_p2, threshold = '0.22mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p2+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp2.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 337.44 mJy
#Peak intensity of source: 28.26 mJy/beam
#rms: 3.93e-02 mJy/beam
#Peak SNR: 718.71
generate_image_png(LB_iteration2_cont_p2+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)


""" Third round of phase-only self-cal """
"""
Check scan length to decide whether to combine scans. If scans are 2 min, do not combine them.
If they are shorter, combine scans here

For PDS 66, scan lenght varies from 2 to 3 minutes for LB, and from 6 to 7 minutes for SB
"""
LB_iteration2_p3 = prefix+'_SBLB_iteration2.p3'
os.system('rm -rf '+LB_iteration2_p3)
gaincal(vis=LB_iteration2_cont_p2+'.ms', caltable=LB_iteration2_p3, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='120s',
        minsnr=2., minblperant=4)
# ~6/47 solutions flagged on average (always under 10)
""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_iteration2_p3,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_p3_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_iteration2_cont_p2+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_p3], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_p3 = prefix+'_SBLB_iteration2_contp3'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p3)
split(vis=LB_iteration2_cont_p2+'.ms', outputvis=LB_iteration2_cont_p3+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_p3+'.ms', imagename = LB_iteration2_cont_p3, threshold = '0.21mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p3+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp3.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 337.09 mJy
#Peak intensity of source: 29.39 mJy/beam
#rms: 3.26e-02 mJy/beam
#Peak SNR: 901.78
generate_image_png(LB_iteration2_cont_p3+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)

""" Fourth round of phase-only self-cal """
LB_iteration2_p4 = prefix+'_SBLB_iteration2.p4'
os.system('rm -rf '+LB_iteration2_p4)
gaincal(vis=LB_iteration2_cont_p3+'.ms', caltable=LB_iteration2_p4, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='60s',
        minsnr=2., minblperant=4)
#~6/47 flagged solution on average (always under 10)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_iteration2_p4,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_p4_gain_phase_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_iteration2_cont_p3+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_p4], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_p4 = prefix+'_SBLB_iteration2_contp4'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p4)
split(vis=LB_iteration2_cont_p3+'.ms', outputvis=LB_iteration2_cont_p4+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_p4+'.ms', imagename = LB_iteration2_cont_p4, threshold = '0.19mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p4+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp4.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 337.09 mJy
#Peak intensity of source: 30.07 mJy/beam
#rms: 3.08e-02 mJy/beam
#Peak SNR: 977.44
generate_image_png(LB_iteration2_cont_p4+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)



""" Fifth round of phase-only self-cal """
LB_iteration2_p5 = prefix+'_SBLB_iteration2.p5'
os.system('rm -rf '+LB_iteration2_p5)
gaincal(vis=LB_iteration2_cont_p4+'.ms', caltable=LB_iteration2_p5, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='30s',
        minsnr=2., minblperant=4)
#~6/47 flagged solution on average for LB (always under 10)

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_iteration2_p5,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_p5_gain_phase_vs_time.png'))

applycal(vis=LB_iteration2_cont_p4+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_p5], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_p5 = prefix+'_SBLB_iteration2_contp5'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p5)
split(vis=LB_iteration2_cont_p4+'.ms', outputvis=LB_iteration2_cont_p5+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_p5+'.ms', imagename = LB_iteration2_cont_p5, threshold = '0.19mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p5+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp5.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 337.16 mJy
#Peak intensity of source: 30.57 mJy/beam
#rms: 2.99e-02 mJy/beam
#Peak SNR: 1024.03
generate_image_png(LB_iteration2_cont_p5+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)


""" Sixth round of phase-only self-cal """
LB_iteration2_p6 = prefix+'_SBLB_iteration2.p6'
os.system('rm -rf '+LB_iteration2_p6)
gaincal(vis=LB_iteration2_cont_p5+'.ms', caltable=LB_iteration2_p6, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='p', solint='18s',
        minsnr=2., minblperant=4)
# up to 12/47 antennas flagged in all LB data

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_p6,xaxis='time', yaxis='GainPhase',iteraxis='spw')
""" Print calibration png file """
plotms(LB_iteration2_p6,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_p6_gain_phase_vs_time.png'))


applycal(vis=LB_iteration2_cont_p5+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_p6], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_p6 = prefix+'_SBLB_iteration2_contp6'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_p6)
split(vis=LB_iteration2_cont_p5+'.ms', outputvis=LB_iteration2_cont_p6+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_p6+'.ms', imagename = LB_iteration2_cont_p6, threshold = '0.18mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p6+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp6.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 337.41 mJy
#Peak intensity of source: 31.11 mJy/beam
#rms: 2.75e-02 mJy/beam
#Peak SNR: 1131.48
generate_image_png(LB_iteration2_cont_p6+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)


""" Clean down to 1 sigma before amplitude self-cal"""

tclean_wrapper(vis=LB_iteration2_cont_p6+'.ms', imagename = LB_iteration2_cont_p6, threshold = '0.028mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_p6+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contp6.image
#Beam 0.138 arcsec x 0.109 arcsec (2.37 deg)
#Flux inside disk mask: 338.86 mJy
#Peak intensity of source: 31.30 mJy/beam
#rms: 2.67e-02 mJy/beam
#Peak SNR: 1171.96



""" Amplitude self-cal"""
LB_iteration2_ap0 = prefix+'_SBLB_iteration2.ap0'
os.system('rm -rf '+LB_iteration2_ap0)
gaincal(vis=LB_iteration2_cont_p6+'.ms', caltable=LB_iteration2_ap0, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw, scan', calmode='ap', solint='inf', minsnr=5.0,
        minblperant=4, solnorm=False)
'''
Only this message
1 of 42 solutions flagged due to SNR < 5 in spw=20 at 2022/07/12/21:52:31.2
'''

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_ap0,xaxis='time', yaxis='GainPhase',iteraxis='spw')
plotms(LB_iteration2_ap0,xaxis='time', yaxis='GainAmp',iteraxis='spw')

""" Print calibration png file """
plotms(LB_iteration2_ap0,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_gain_ap0_phase_vs_time.png'))
plotms(LB_iteration2_ap0,xaxis='time', yaxis='GainAmp',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_gain_ap0_amp_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_iteration2_cont_p6+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_ap0], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_ap0 = prefix+'_SBLB_iteration2_contap0'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_ap0)
split(vis=LB_iteration2_cont_p6+'.ms', outputvis=LB_iteration2_cont_ap0+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
#clean again down to 1 sigma
tclean_wrapper(vis=LB_iteration2_cont_ap0+'.ms', imagename = LB_iteration2_cont_ap0, threshold = '0.027mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_ap0+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contap0.image
#Beam 0.140 arcsec x 0.110 arcsec (5.22 deg)
#Flux inside disk mask: 339.44 mJy
#Peak intensity of source: 32.22 mJy/beam
#rms: 2.62e-02 mJy/beam
#Peak SNR: 1228.98
generate_image_png(LB_iteration2_cont_ap0+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)


""" Try ampl self-cal on scan length intervals"""
LB_iteration2_ap1 = prefix+'_SBLB_iteration2.ap1'
os.system('rm -rf '+LB_iteration2_ap1)
gaincal(vis=LB_iteration2_cont_ap0+'.ms', caltable=LB_iteration2_ap1, gaintype='T', spw=LB_contspws,
        refant=LB_refant, combine='spw', calmode='ap', solint='inf', minsnr=5.0,
        minblperant=4, solnorm=False)
#~6/47 flagged solution on average (always under 10) 

""" Inspect gain tables interactively and decide whether to manually flag something"""
plotms(LB_iteration2_ap1,xaxis='time', yaxis='GainPhase',iteraxis='spw')
plotms(LB_iteration2_ap1,xaxis='time', yaxis='GainAmp',iteraxis='spw')

""" Print calibration png file """
plotms(LB_iteration2_ap1,xaxis='time', yaxis='GainPhase',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_gain_ap1_phase_vs_time.png'))
plotms(LB_iteration2_ap1,xaxis='time', yaxis='GainAmp',iteraxis='spw',exprange='all',
       overwrite=True,showgui=False,
       plotfile=os.path.join(LB_selfcal_iteration2_folder,f'{prefix}_LB_iteration2_gain_ap1_amp_vs_time.png'))

""" Apply the solutions """
applycal(vis=LB_iteration2_cont_ap0+'.ms', spw=LB_contspws, spwmap = LB_spw_mapping,
         gaintable=[LB_iteration2_ap1], interp='linearPD', calwt=True, applymode='calonly')

""" Split off a corrected MS """
LB_iteration2_cont_ap1 = prefix+'_SBLB_iteration2_contap1'
os.system('rm -rf %s.ms*' % LB_iteration2_cont_ap1)
split(vis=LB_iteration2_cont_ap0+'.ms', outputvis=LB_iteration2_cont_ap1+'.ms', datacolumn='corrected')

""" Image the results; check the resulting map """
tclean_wrapper(vis=LB_iteration2_cont_ap1+'.ms', imagename = LB_iteration2_cont_ap1,threshold = '0.027mJy',
               **LB_tclean_wrapper_kwargs)
estimate_SNR(LB_iteration2_cont_ap1+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
#PDS_66_SBLB_iteration2_contap1.image
#Beam 0.140 arcsec x 0.110 arcsec (5.41 deg)
#Flux inside disk mask: 339.76 mJy
#Peak intensity of source: 32.05 mJy/beam
#rms: 2.58e-02 mJy/beam
#Peak SNR: 1239.94
generate_image_png(LB_iteration2_cont_ap1+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_iteration2_LB,10*rms_iteration2_LB],
                   save_folder=LB_selfcal_iteration2_folder)

# only minor increase of SNR, but images look fine


#Now check again the flux scaling


self_caled_LB_iteration2_visibilities = {'p1':LB_iteration2_cont_p1,
                                         'p2':LB_iteration2_cont_p2,
                                         'p3':LB_iteration2_cont_p3,
                                         'p4':LB_iteration2_cont_p4,
                                         'p5':LB_iteration2_cont_p5,
                                         'p6':LB_iteration2_cont_p6,
                                         'ap0':LB_iteration2_cont_ap0,
                                         'ap1':LB_iteration2_cont_ap1}


for vis in self_caled_LB_iteration2_visibilities.values(): 
    listobs(vis=vis+'.ms',listfile=vis+'.ms.txt',overwrite=True)

LB_EBs = ('EB0','EB1','EB2','EB3')
LB_EB_spws = ('0,1,2,3', '4,5,6,7', '8,9,10,11', '20,21,22,23') #fill out by referring to listobs output

for self_cal_step,self_caled_vis in self_caled_LB_iteration2_visibilities.items():
    for EB_key,spw in zip(LB_EBs,LB_EB_spws):
        nametemplate = f'{prefix}_LB_iteration2_{EB_key}_{self_cal_step}_compare_amp_vs_time'
        visibilities = [self_caled_vis+'.ms',LB_iteration2_cont_p0+'.ms']
        plot_amp_vs_time_comparison(
                nametemplate=nametemplate,visibilities=visibilities,spw=spw,
                uvrange=uv_ranges['LB'],output_folder=LB_selfcal_iteration2_folder)

#set to the EB of the combined SBLB data that corresponds to flux_ref_EB
SBLB_flux_ref_EB = 1 #this is LB EB1

total_number_of_EBs = number_of_EBs['SB'] + number_of_EBs['LB']
for self_cal_step,vis in self_caled_LB_iteration2_visibilities.items():
    nametemplate = f'{prefix}_SBLB_iteration2_cont{self_cal_step}_EB'
    split_all_obs(msfile=vis+'.ms',nametemplate=nametemplate)
    for i in range(total_number_of_EBs):
        export_MS(f'{nametemplate}{i}.ms')
    for i in range(total_number_of_EBs):
        reference = f'{nametemplate}{SBLB_flux_ref_EB}.vis.npz'
        output = f'iteration2_flux_comparison_EB{i}_to_EB{SBLB_flux_ref_EB}'\
                       +f'_SBLB_{self_cal_step}.png'
        plot_label = os.path.join(LB_selfcal_iteration2_folder,output)
        estimate_flux_scale(reference=reference,
                            comparison=f'{nametemplate}{i}.vis.npz',
                            incl=incl, PA=PA,
                            plot_label=plot_label)
    filelist = [f'{nametemplate}{i}.vis.npz' for i in range(total_number_of_EBs)]
    fluxscale = [1.,]*total_number_of_EBs
    plot_label = os.path.join(LB_selfcal_iteration2_folder,
                              f'deprojected_vis_profiles_SBLB_iteration2_{self_cal_step}.png')
    plot_deprojected(filelist=filelist,fluxscale=fluxscale, PA=PA, incl=incl,
                     show_err=True,plot_label=plot_label)

#The ratio of the fluxes of PDS_66_SBLB_iteration2_contap1_EB0.vis.npz to
#PDS_66_SBLB_iteration2_contap1_EB1.vis.npz is 0.99863
#The scaling factor for gencal is 0.999 for your comparison measurement
#The error on the weighted mean ratio is 2.935e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_iteration2_contap1_EB1.vis.npz to
#PDS_66_SBLB_iteration2_contap1_EB1.vis.npz is 1.00000
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 3.025e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_iteration2_contap1_EB2.vis.npz to
#PDS_66_SBLB_iteration2_contap1_EB1.vis.npz is 0.99972
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 3.668e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_iteration2_contap1_EB3.vis.npz to
#PDS_66_SBLB_iteration2_contap1_EB1.vis.npz is 0.99992
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 2.981e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_iteration2_contap1_EB4.vis.npz to
#PDS_66_SBLB_iteration2_contap1_EB1.vis.npz is 1.00067
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 2.735e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor

#The ratio of the fluxes of PDS_66_SBLB_iteration2_contap1_EB5.vis.npz to
#PDS_66_SBLB_iteration2_contap1_EB1.vis.npz is 1.00055
#The scaling factor for gencal is 1.000 for your comparison measurement
#The error on the weighted mean ratio is 2.640e-04, although it's likely that
#the weights in the measurement sets are too off by some constant factor


'''
Comments: 

The procedure with two iterations of both SB and SB+LB phase self-cal worked perfectly!
Now, there no more waterfall features in the amplitudes nor downward trends in the flux ratio plots for both SB and LB EBs.
The flux ratios for all EBs after this procedure are well below the 4% treshold.
'''

############ END of COMBINED SB+LB phase-only self-cal iteration 2 ############


""" Split out final continuum ms table, with a 30s timebin
"""
#comparing ap0 and ap1, the SNR is almost the same, and the flux scaling is fine for both
#however, ap1 seems to have a slightly "better" noise structure when looking at the
#images. So I choose ap1.

LB_iteration2_cont_averaged = f'{prefix}_time_ave_continuum'
os.system(f'rm -rf {LB_iteration2_cont_averaged}.ms*')
split(vis=LB_iteration2_cont_ap1+'.ms', outputvis=LB_iteration2_cont_averaged+'.ms', datacolumn='data',
      keepflags=False, timebin='30s')


"""
Now apply these solutions to the line data
"""

calibrate_linedata_folder = get_figures_folderpath('9_apply_cal_to_lines')
make_figures_folder(calibrate_linedata_folder)

""" Check that lines are not flagged in not averaged data"""
for params in data_params.values():
    plotfile = os.path.join(
                calibrate_linedata_folder,
                f'{prefix}_{params["name"]}_chan-v-amp_preselfcal_after_flagging.png')
    plotms(vis=params['vis'],
           xaxis='channel',
           yaxis='amplitude',
           field=params['field'],
           ydatacolumn='data',
           avgtime='1e8',
           avgscan=True,
           avgbaseline=True,
           coloraxis='corr',
           iteraxis='spw',
           showgui = False,
           exprange='all',
           plotfile=plotfile)


""" Apply gaintables of individual EBs"""
for params in data_params.values():
    single_EB_p1 = f'{prefix}_{params["name"]}_initcont.p1'
    vis = f'{prefix}_{params["name"]}.ms'
    applycal(vis=vis,spw='0~3',spwmap=[0,0,0,0],gaintable=[single_EB_p1],
             interp='linearPD',applymode='calonly',calwt=True)
    split(vis=vis,outputvis=vis[:-3]+'_no_ave_selfcal.ms',datacolumn='corrected')

### ALIGN DATA ###
# we re-align the no_ave data, as we have done for the "initcont" ms tables.
# We go from *_no_ave_selfcal.ms to *_no_ave_shift.ms

print('Align no_ave data')

reference_ms = {'LB':reference_for_LB_alignment,
                'SB':reference_for_SB_alignment}
for params in data_params.values():
    unshifted_ms = f'{prefix}_{params["name"]}_no_ave_selfcal.ms'
    array_key,_ = params['name'].split('_') #LB or SB
    offset = alignment_offsets[params['name']]
    #npix and cell_size are not needed because we do not fit any offset
    alignment.align_measurement_sets(
            reference_ms=reference_ms[array_key],align_ms=[unshifted_ms],
            align_offsets=[offset],npix=None,cell_size=None)

### END OF ALING DATA ###


"""
If you have re-scaled fluxes, you need to re-scale the shifted *no_ave* EBs as well
"""

rescale_flux(vis=prefix+'_SB_EB0_no_ave_selfcal_shift.ms', gencalparameter=[0.948])
rescale_flux(vis=prefix+'_SB_EB1_no_ave_selfcal_shift.ms', gencalparameter=[0.962])
rescale_flux(vis=prefix+'_LB_EB0_no_ave_selfcal_shift.ms', gencalparameter=[1.020])
rescale_flux(vis=prefix+'_LB_EB3_no_ave_selfcal_shift.ms', gencalparameter=[0.912])

listobs(vis=prefix+'_SB_EB0_no_ave_selfcal_shift_rescaled.ms',listfile=prefix+'_SB_EB0_no_ave_selfcal_shift_rescaled.ms'+'.txt',overwrite=True)
listobs(vis=prefix+'_SB_EB1_no_ave_selfcal_shift_rescaled.ms',listfile=prefix+'_SB_EB1_no_ave_selfcal_shift_rescaled.ms'+'.txt',overwrite=True)
listobs(vis=prefix+'_LB_EB0_no_ave_selfcal_shift_rescaled.ms',listfile=prefix+'_LB_EB0_no_ave_selfcal_shift_rescaled.ms'+'.txt',overwrite=True)
listobs(vis=prefix+'_LB_EB3_no_ave_selfcal_shift_rescaled.ms',listfile=prefix+'_LB_EB3_no_ave_selfcal_shift_rescaled.ms'+'.txt',overwrite=True)



""" Concat not averaged SB data """
SB_combined = f'{prefix}_SB_no_ave_concat'
os.system('rm -rf %s.ms*' % SB_combined)
concat(vis=[f'{prefix}_SB_EB0_no_ave_selfcal_shift_rescaled.ms',
            f'{prefix}_SB_EB1_no_ave_selfcal_shift_rescaled.ms'],
       concatvis=SB_combined+'.ms', dirtol='0.1arcsec', copypointing=False)

listobs(vis=SB_combined+'.ms',listfile=SB_combined+'.ms.txt',overwrite=True)

# BE CAREFUL HERE
# using gaintables from iteration2
applycal(vis=SB_combined+'.ms', spw='0~7',
         gaintable=[SB_iteration2_p1,SB_iteration2_p2,SB_iteration2_p3,SB_iteration2_p4,SB_iteration2_p5],
         spwmap = [SB_spw_mapping]*5,interp=['linearPD']*5, calwt=True,
         applymode='calonly',flagbackup=False)

SB_no_ave_selfcal = f'{prefix}_SB_no_ave_selfcal.ms'
os.system(f'rm -rf {SB_no_ave_selfcal}*')
split(vis=SB_combined+'.ms', outputvis=SB_no_ave_selfcal,datacolumn='corrected')

""" Concat not averaged LB data """

LB_combined = f'{prefix}_SBLB_no_ave_concat.ms'
os.system(f'rm -rf {LB_combined}*')

concat(vis=[SB_no_ave_selfcal]+[f'{prefix}_LB_EB0_no_ave_selfcal_shift_rescaled.ms',
                                f'{prefix}_LB_EB1_no_ave_selfcal_shift.ms',
                                f'{prefix}_LB_EB2_no_ave_selfcal_shift.ms',
                                f'{prefix}_LB_EB3_no_ave_selfcal_shift_rescaled.ms',],
       concatvis=LB_combined, dirtol='0.1arcsec', copypointing=False)

listobs(vis=LB_combined,listfile=LB_combined+'.txt',overwrite=True)

# BE CAREFUL HERE
# using all gaintables from iteration2 even for LB
applycal(vis=LB_combined, spw='0~23',
         gaintable=[LB_iteration2_p1,LB_iteration2_p2,LB_iteration2_p3,LB_iteration2_p4,LB_iteration2_p5,LB_iteration2_p6,
                    LB_iteration2_ap0,LB_iteration2_ap1],
         spwmap = [LB_spw_mapping]*8,interp=['linearPD']*8,
         calwt=True, applymode='calonly',flagbackup=False)

SBLB_no_ave_selfcal = f'{prefix}_SBLB_no_ave_selfcal_time_ave.ms'
os.system(f'rm -rf {SBLB_no_ave_selfcal}*')
""" time average of 30s, tests show there is no difference with data without time average """
split(vis=LB_combined, outputvis=SBLB_no_ave_selfcal,
      datacolumn='corrected', timebin = '30s', keepflags=False)

listobs(vis=SBLB_no_ave_selfcal,listfile=SBLB_no_ave_selfcal+'.txt',overwrite=True)

# Check that solutions have been applied correctly by flagging the line data, averaging and imaging continuum
# continuum has to be the same imaged in the last step of the self-cal
complete_dataset_dict = {'vis' : SBLB_no_ave_selfcal,
                         'name' : 'SBLB_concat',
                         'field' : fields['LB'],
                         'line_spws': np.array([0,2,3,
                                                4,6,7,
                                                8,10,11,
                                                12,14,15,
                                                16,18,19,
                                                20,22,23]), # list of spws containing lines
                         'line_freqs': np.array([rest_freq_13CO,rest_freq_CS,rest_freq_12CO]*6), #frequencies (Hz) corresponding to line_spws
                         'cont_spws': None,
                         'width_array': None,
                         }

#use the output of get_flagchannels at the beginning of the script to define fitspw

# Flagchannels input string for LB_EB0: '0:836~3004, 2:794~3043, 3:787~3055'
# Flagchannels input string for LB_EB1: '0:837~3005, 2:794~3043, 3:787~3055'
# Flagchannels input string for LB_EB2: '0:837~3005, 2:794~3043, 3:787~3055'
# Flagchannels input string for LB_EB3: '0:835~3003, 2:796~3045, 3:784~3051'
# Flagchannels input string for SB_EB0: '0:837~3005, 2:794~3043, 3:788~3056'
# Flagchannels input string for SB_EB1: '0:837~3005, 2:794~3043, 3:788~3056'

# remember how spws are organized in the SBLB files (seen from listobs)
# [0,1,2,3]     -> LB EB0
# [4,5,6,7]     -> LB EB1
# [8,9,10,11]   -> LB EB2
# [12,13,14,15] -> SB EB0
# [16,17,18,19] -> SB EB1
# [20,21,22,23] -> LB EB3
fitspw =  '0:836~3004, 1:0, 2:794~3043, 3:787~3055,'\
         +'4:837~3005, 5:0, 6:794~3043, 7:787~3055,'\
         +'8:837~3005, 9:0, 10:794~3043, 11:787~3055,'\
         +'12:837~3005, 13:0, 14:794~3043, 15:788~3056,'\
         +'16:837~3005, 17:0, 18:794~3043, 19:788~3056,'\
         +'20:835~3003, 21:0, 22:796~3045, 23:784~3051'

avg_cont(ms_dict=complete_dataset_dict, output_prefix=prefix, flagchannels = fitspw,
         contspws = complete_dataset_dict['cont_spws'], width_array=complete_dataset_dict['width_array'])

# image the avg then cal with same parameters as in last step of self-cal
tclean_wrapper(vis=LB_iteration2_cont_averaged+'.ms',
               imagename = LB_iteration2_cont_averaged+'_image', mask=LB_mask,
               threshold = '0.027mJy', deconvolver='multiscale', scales=scales['LB'],
               imsize=imsize['LB'], cellsize=cellsize['LB'],
               robust=0.5, interactive=False, parallel=use_parallel, gridder='standard')
estimate_SNR(LB_iteration2_cont_averaged+'_image'+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
generate_image_png(LB_iteration2_cont_averaged+'_image'+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],save_folder=calibrate_linedata_folder)
#PDS_66_time_ave_continuum_image.image
#Beam 0.140 arcsec x 0.110 arcsec (5.41 deg)
#Flux inside disk mask: 339.74 mJy
#Peak intensity of source: 32.06 mJy/beam
#rms: 2.58e-02 mJy/beam
#Peak SNR: 1240.73


# image the cal then avg with same parameters as in last step of self-cal
complete_dataset_image = prefix+'_'+complete_dataset_dict['name']+'_initcont_image'
tclean_wrapper(vis=prefix+'_'+complete_dataset_dict['name']+'_initcont.ms',
               imagename = complete_dataset_image, mask=LB_mask,
               threshold = '0.027mJy', deconvolver='multiscale', scales=scales['LB'],
               imsize=imsize['LB'], cellsize=cellsize['LB'],
               robust=0.5, interactive=False, parallel=use_parallel, gridder='standard')
estimate_SNR(complete_dataset_image+'.image', disk_mask = LB_mask, noise_mask = noise_annulus_LB)
generate_image_png(complete_dataset_image+'.image',plot_sizes=image_png_plot_sizes,
                   color_scale_limits=[-3*rms_LB,10*rms_LB],save_folder=calibrate_linedata_folder)
#PDS_66_SBLB_concat_initcont_image.image
#Beam 0.140 arcsec x 0.110 arcsec (5.41 deg)
#Flux inside disk mask: 339.73 mJy
#Peak intensity of source: 32.06 mJy/beam
#rms: 2.58e-02 mJy/beam
#Peak SNR: 1240.77

# Plot ratio of cal then avg, and avg then cal. It should be equal to ~one (only difference being the time average):
ref_image = LB_iteration2_cont_averaged+'_image'+'.image'
os.system('rm -rf '+complete_dataset_image+'.ratio')
immath(imagename=[ref_image,complete_dataset_image+'.image'],mode='evalexpr',
       outfile=complete_dataset_image+'.ratio',
       expr='iif(IM0 > 3*'+str(rms_LB)+', IM1/IM0, 0)'
       )
generate_image_png(f'{complete_dataset_image}.ratio',plot_sizes=[2*mask_semimajor,2*mask_semimajor],
                   color_scale_limits=[0.5,1.5],image_units='ratio',
                   save_folder=calibrate_linedata_folder)
# Check OK, the ratio is equal to one

# Do continuum subtraction
contsub_vis = f'{SBLB_no_ave_selfcal}.contsub'
os.system(f'rm -rf {contsub_vis}*')
uvcontsub(vis=SBLB_no_ave_selfcal, spw='0~23', fitspw=fitspw,
          excludechans=True, solint='int', fitorder=1, want_cont=False)
'''
2022-11-04 08:45:23     WARN    VBContinuumSubtractor::apply    Extrapolating to cover [330.624, 332.498] (GHz).
2022-11-04 08:45:23     WARN    VBContinuumSubtractor::apply+   The frequency range used for the continuum fit was [330.624, 332.498] (GHz).
2022-11-04 08:48:46     WARN    VBContinuumSubtractor::apply    Extrapolating to cover [330.625, 332.5] (GHz).
2022-11-04 08:48:46     WARN    VBContinuumSubtractor::apply+   The frequency range used for the continuum fit was [330.625, 332.499] (GHz).
2022-11-04 08:51:57     WARN    VBContinuumSubtractor::apply    Extrapolating to cover [330.625, 332.5] (GHz).
2022-11-04 08:51:57     WARN    VBContinuumSubtractor::apply+   The frequency range used for the continuum fit was [330.625, 332.499] (GHz).
2022-11-04 08:55:15     WARN    VBContinuumSubtractor::apply    Extrapolating to cover [330.63, 332.505] (GHz).
2022-11-04 08:55:15     WARN    VBContinuumSubtractor::apply+   The frequency range used for the continuum fit was [330.63, 332.504] (GHz).
2022-11-04 08:56:21     WARN    VBContinuumSubtractor::apply    Extrapolating to cover [330.605, 332.479] (GHz).
2022-11-04 08:56:21     WARN    VBContinuumSubtractor::apply+   The frequency range used for the continuum fit was [330.605, 332.479] (GHz).
2022-11-04 08:58:00     WARN    VBContinuumSubtractor::apply    Extrapolating to cover [330.63, 332.505] (GHz).
2022-11-04 08:58:00     WARN    VBContinuumSubtractor::apply+   The frequency range used for the continuum fit was [330.63, 332.504] (GHz).
'''

listobs(vis=SBLB_no_ave_selfcal+'.contsub',listfile=SBLB_no_ave_selfcal+'.contsub.txt',overwrite=True)

# Split final ms table into separate spws for 12CO, 13CO, CS and continuum

#12CO
vis_12CO = SBLB_no_ave_selfcal[:-3]+'_12CO.ms'
os.system(f'rm -rf {vis_12CO}*')
spw_12CO = '3,7,11,15,19,23'

split(vis=SBLB_no_ave_selfcal,outputvis=vis_12CO,spw=spw_12CO,
      datacolumn='data', keepflags=False)
listobs(vis=vis_12CO,listfile=vis_12CO+'.txt',overwrite=True)

split(vis=contsub_vis,outputvis=f'{vis_12CO}.contsub',spw=spw_12CO,
      datacolumn='data', keepflags=False)
listobs(vis=vis_12CO+'.contsub',listfile=vis_12CO+'.contsub.txt',overwrite=True)

#13CO
vis_13CO = SBLB_no_ave_selfcal[:-3]+'_13CO.ms'
os.system(f'rm -rf {vis_13CO}*')
spw_13CO = '0,4,8,12,16,20'

split(vis=SBLB_no_ave_selfcal,outputvis=vis_13CO,spw=spw_13CO,
      datacolumn='data', keepflags=False)
listobs(vis=vis_13CO,listfile=vis_13CO+'.txt',overwrite=True)

split(vis=contsub_vis,outputvis=f'{vis_13CO}.contsub',spw=spw_13CO,
      datacolumn='data', keepflags=False)
listobs(vis=vis_13CO+'.contsub',listfile=vis_13CO+'.contsub.txt',overwrite=True)

#CS
vis_CS = SBLB_no_ave_selfcal[:-3]+'_CS.ms'
os.system(f'rm -rf {vis_CS}*')
spw_CS = '2,6,10,14,18,22'

split(vis=SBLB_no_ave_selfcal,outputvis=vis_CS,spw=spw_CS,
      datacolumn='data', keepflags=False)
listobs(vis=vis_CS,listfile=vis_CS+'.txt',overwrite=True)

split(vis=contsub_vis,outputvis=f'{vis_CS}.contsub',spw=spw_CS,
      datacolumn='data', keepflags=False)
listobs(vis=vis_CS+'.contsub',listfile=vis_CS+'.contsub.txt',overwrite=True)

#continuum spw
vis_continuum = SBLB_no_ave_selfcal[:-3]+'_contspw.ms'
os.system(f'rm -rf {vis_continuum}*')
spw_continuum = '1,5,9,13,17,21'

split(vis=SBLB_no_ave_selfcal,outputvis=vis_continuum,spw=spw_continuum,
      datacolumn='data', keepflags=False)
listobs(vis=vis_continuum,listfile=vis_continuum+'.txt',overwrite=True)

split(vis=contsub_vis,outputvis=f'{vis_continuum}.contsub',spw=spw_continuum,
      datacolumn='data', keepflags=False)
listobs(vis=vis_continuum+'.contsub',listfile=vis_continuum+'.contsub.txt',overwrite=True)
